🎨 update type hints.

cemc-oper · Jun 20, 2024 · b912561 · b912561
1 parent 93a54be
commit b912561
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 7 deletions.
diff --git a/reki/format/grib/__init__.py b/reki/format/grib/__init__.py
@@ -1,4 +1,4 @@
-from typing import Union, Optional, Dict
+from typing import Union, Optional, Dict, Literal
 from pathlib import Path
 
 import xarray as xr
@@ -12,7 +12,7 @@ def load_field_from_file(
         parameter: Union[str, Dict],
         level_type: Optional[Union[str, Dict]] = None,
         level: Optional[int] = None,
-        engine: str = "eccodes",
+        engine: Literal["eccodes", "cfgrib"] = "eccodes",
         **kwargs
 ) -> Optional[xr.DataArray]:
     """
@@ -36,7 +36,7 @@ def load_field_from_file(
 
     Returns
     -------
-    Union[DataArray]
+    Optional[xr.DataArray]
         DataArray if found one field, or None if not.
     """
     fixed_level_type = fix_level_type(level_type, engine=engine)

diff --git a/reki/format/grib/common/_level.py b/reki/format/grib/common/_level.py
@@ -1,9 +1,9 @@
-from typing import Optional, Union, Dict
+from typing import Optional, Union, Dict, Literal
 
 
 def fix_level_type(
         level_type: Optional[Union[str, Dict]],
-        engine: str = "eccodes",
+        engine: Literal["eccodes", "cfgrib"] = "eccodes",
 ) -> Optional[Union[str, Dict]]:
     """
     Convert level type into dict of GRIB key(s).
@@ -23,7 +23,7 @@ def fix_level_type(
     Parameters
     ----------
     level_type
-    engine : str
+    engine
         load engine:
 
         - cfgrib: use cfgrib to load message

diff --git a/reki/format/netcdf/__init__.py b/reki/format/netcdf/__init__.py
@@ -1,5 +1,4 @@
 from pathlib import Path
-import typing
 from typing import Union, Dict, Optional
 
 import xarray as xr