Add type hints to public API

Following the recommendation from numpy/numpy#24300

src/ecl_data_io/array_entry.py

@@ -1,4 +1,10 @@
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional, Union
+
+if TYPE_CHECKING:
+    from numpy.typing import ArrayLike
+
+    from .types import ArrayValue
 
 
 class EclArray(ABC):
@@ -35,45 +41,49 @@ def is_eof(self):
             self._read()
         return self._is_eof
 
-    def read_keyword(self):
+    def read_keyword(self) -> str:
         """
         Read the keyword from the ecl file.
 
         :returns: The keyword as a 8 character string.
         """
         if self._keyword is None:
             self._read()
-        return self._keyword
+        return self._keyword  # type: ignore
 
-    def read_length(self):
+    def read_length(self) -> int:
         """
         Read the length from the ecl file.
 
         :returns: The length of the array in number of entries.
         """
         if self._length is None:
             self._read()
-        return self._length
+        return self._length  # type: ignore
 
     @abstractmethod
-    def read_array(self):
+    def read_array(self) -> "ArrayValue":
         """
         Read the array from the unformatted ecl file.
 
         :returns: numpy array of values.
         """
         pass
 
-    def read_type(self):
+    def read_type(self) -> Union[str, bytes]:
         """
-        The type given in the header of the array
+        The type given in the header of the array.
+        In case of unformatted file this will be bytes, for
+        unformatted files it is str.
         """
         if self._type is None:
             self._read()
-        return self._type
+        return self._type  # type: ignore
 
     @abstractmethod
-    def update(self, *, keyword=None, array=None):
+    def update(
+        self, *, keyword: Optional[str] = None, array: Optional["ArrayLike"] = None
+    ):
         """
         Updates the entry with the given new data.
 

src/ecl_data_io/format.py

@@ -41,7 +41,7 @@ def guess_format(filelike):
         return Format.FORMATTED
 
 
-def get_stream(filepath, fileformat, mode="r"):
+def get_stream(filepath, fileformat: Format, mode: str = "r"):
     """
     Openes the given file with the correct mode (text or binary)
     based on fileformat.
@@ -58,7 +58,7 @@ def get_stream(filepath, fileformat, mode="r"):
         return filepath, False
 
 
-def check_correct_mode(stream, fileformat):
+def check_correct_mode(stream, fileformat: Format):
     """
     Checks that the stream is the correct mode (text or binary) for the given
     fileformat.

src/ecl_data_io/read.py

@@ -1,11 +1,15 @@
-from pathlib import Path
+from typing import TYPE_CHECKING, Iterator, List, Optional, Tuple
 
 from ecl_data_io._formatted.read import FormattedEclArray
 from ecl_data_io._unformatted.read import UnformattedEclArray
+from ecl_data_io.array_entry import EclArray
 from ecl_data_io.format import Format, check_correct_mode, get_stream, guess_format
 
+if TYPE_CHECKING:
+    from .types import ArrayValue
 
-def read(*args, **kwargs):
+
+def read(*args, **kwargs) -> List[Tuple[str, "ArrayValue"]]:
     """
     Read the contents of a ecl file and return a list of
     tuples (keyword, array). Takes the same parameters as
@@ -16,7 +20,7 @@ def read(*args, **kwargs):
     ]
 
 
-def lazy_read(filelike, fileformat=None):
+def lazy_read(filelike, fileformat: Optional[Format] = None) -> Iterator[EclArray]:
     """
     Reads the contents of an ecl file and generates the entries
     of that file. Each entry has a entry.read_keyword() and

src/ecl_data_io/types.py

@@ -23,8 +23,10 @@
 
 """
 import warnings
+from typing import Union
 
 import numpy as np
+from numpy.typing import ArrayLike
 
 # np dtype for ecl types with fixed width
 static_dtypes = {
@@ -45,6 +47,9 @@ class MESS:
     pass
 
 
+ArrayValue = Union[ArrayLike, MESS]
+
+
 def to_np_type(type_keyword):
     """
     :param type_keyword: A bytestring of a ecl type.

src/ecl_data_io/write.py

@@ -1,11 +1,17 @@
-from pathlib import Path
+from typing import Dict, Sequence, Tuple, Union
 
 from ecl_data_io._formatted.write import formatted_write
 from ecl_data_io._unformatted.write import unformatted_write
 from ecl_data_io.format import Format, check_correct_mode, get_stream
 
+from .types import ArrayValue
 
-def write(filelike, contents, fileformat=Format.UNFORMATTED):
+
+def write(
+    filelike,
+    contents: Union[Sequence[Tuple[str, ArrayValue]], Dict[str, ArrayValue]],
+    fileformat: Format = Format.UNFORMATTED,
+):
     """
     Write the given contents to the given file in ecl format.