CLN: Add types to ecl_output_file

mferrera · Nov 23, 2023 · 434097a · 434097a
1 parent 25868bd
commit 434097a
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/xtgeo/grid3d/_ecl_output_file.py b/src/xtgeo/grid3d/_ecl_output_file.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 from enum import Enum, unique
+from typing import Literal
 
 
 @unique
@@ -15,7 +18,7 @@ class TypeOfGrid(Enum):
     BLOCK_CENTER = 3
 
     @classmethod
-    def alternate_code(cls, code):
+    def alternate_code(cls, code: Literal[0, 1, 2, 3]) -> TypeOfGrid:
         """Converts from alternate code to TypeOfGrid member.
 
         weirdly, type of grid sometimes (For instance init's INTHEAD and
@@ -34,7 +37,7 @@ def alternate_code(cls, code):
         return type_of_grid
 
     @property
-    def alternate_value(self):
+    def alternate_value(self) -> int:
         """Inverse of alternate_code."""
         alternate_value = 0
         if self == TypeOfGrid.CORNER_POINT: