Add type annotations to forward_models

Author: einarwar

pyproject.toml

@@ -132,7 +132,7 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"!*/semeio/fmudesign/**.py" = ["ANN"]
+"!*/semeio/{fmudesign,forward_models}/**.py" = ["ANN"]
 
 [tool.ruff.lint.pylint]
 max-args = 20

src/semeio/forward_models/__init__.py

@@ -13,7 +13,7 @@
 
 
 class Design2Params(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="DESIGN2PARAMS",
             command=[
@@ -44,7 +44,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class DesignKW(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="DESIGN_KW",
             command=[
@@ -86,7 +86,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class GenDataRFT(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="GENDATA_RFT",
             command=[
@@ -160,7 +160,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class OTS(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="OTS",
             command=["overburden_timeshift", "-c", "<CONFIG>"],
@@ -180,7 +180,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class Pyscal(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="PYSCAL",
             command=[
@@ -223,7 +223,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class InsertNoSim(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="INSERT_NOSIM",
             command=[
@@ -255,7 +255,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class RemoveNoSim(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="REMOVE_NOSIM",
             command=["sed", "-i", "", "/^NOSIM/d", "<ECLBASE>.DATA"]
@@ -277,7 +277,7 @@ def documentation() -> ForwardModelStepDocumentation | None:
 
 
 class ReplaceString(ForwardModelStepPlugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             name="REPLACE_STRING",
             command=["replace_string", "-o", "<FROM>", "-n", "<TO>", "-f", "<FILE>"],

src/semeio/forward_models/design2params/design2params.py

@@ -1,6 +1,8 @@
 import logging
 import warnings
+from collections.abc import Sequence
 from pathlib import Path
+from typing import Literal
 
 import numpy as np
 import pandas as pd
@@ -24,13 +26,13 @@
 
 
 def run(
-    realization,
-    xlsfilename,
-    designsheetname="DesignSheet01",
-    defaultssheetname="DefaultValues",
-    parametersfilename="parameters.txt",
-    log_level=None,
-):
+    realization: int,
+    xlsfilename: str,
+    designsheetname: str = "DesignSheet01",
+    defaultssheetname: str = "DefaultValues",
+    parametersfilename: str = "parameters.txt",
+    log_level: int | str | None = None,
+) -> None:
     """
     Reads out all file content from different files and create dataframes
     """
@@ -68,8 +70,12 @@ def run(
 
 
 def _complete_parameters_file(
-    realization, parameters, parametersfilename, design_matrix_sheet, default_sheet
-):
+    realization: int,
+    parameters: pd.DataFrame,
+    parametersfilename: str,
+    design_matrix_sheet: pd.DataFrame,
+    default_sheet: pd.DataFrame,
+) -> None:
     """
     Pick key / values from chosen realization in design matrix
     Append those key / values if not present into parameters.txt
@@ -184,7 +190,13 @@ def _complete_parameters_file(
     )
 
 
-def _read_excel(file_name, sheet_name, header=0, usecols=None, engine=None):
+def _read_excel(
+    file_name: str,
+    sheet_name: str,
+    header: int | Sequence[int] | None = 0,
+    usecols: list[int] | None = None,
+    engine: Literal["xlrd", "openpyxl", "odf", "pyxlsb", "calamine"] | None = None,
+) -> pd.DataFrame:
     """
     Make dataframe from excel file
     :return: Dataframe
@@ -219,7 +231,7 @@ def _read_excel(file_name, sheet_name, header=0, usecols=None, engine=None):
     return dframe.dropna(axis=1, how="all")
 
 
-def _validate_design_matrix_header(design_matrix):
+def _validate_design_matrix_header(design_matrix: pd.DataFrame) -> None:
     """
     Validate header in user inputted design matrix
     :raises: ValueError if design matrix contains empty headers
@@ -239,7 +251,7 @@ def _validate_design_matrix_header(design_matrix):
         raise ValueError(f"Column headers not present in column {column_indexes}")
 
 
-def _invalid_design_realizations(design_matrix):
+def _invalid_design_realizations(design_matrix: pd.DataFrame) -> set[int]:
     """
     Build a set of realization indices where something is wrong,
     f.ex empty cells
@@ -263,7 +275,7 @@ def _invalid_design_realizations(design_matrix):
 
     # Look for initial or trailing whitespace in column headers. This
     # is disallowed as it can create user confusion and has no use-case.
-    for col_header in design_matrix:
+    for col_header in design_matrix.columns:
         if col_header != col_header.strip():
             raise SystemExit(
                 f'Column header "{col_header}" contains initial or trailing whitespace.'
@@ -292,7 +304,7 @@ def _invalid_design_realizations(design_matrix):
     return {cell_coord[0] for cell_coord in empty_cell_coords}
 
 
-def _read_defaultssheet(xlsfilename, defaultssheetname):
+def _read_defaultssheet(xlsfilename: str, defaultssheetname: str) -> pd.DataFrame:
     """
     Construct a dataframe of keys and values to be used as defaults from the
     first two columns in a spreadsheet.
@@ -304,7 +316,10 @@ def _read_defaultssheet(xlsfilename, defaultssheetname):
     """
     if defaultssheetname:
         default_df = _read_excel(
-            xlsfilename, defaultssheetname, usecols=[0, 1], header=None
+            xlsfilename,
+            defaultssheetname,
+            usecols=[0, 1],
+            header=None,
         )
         if default_df.empty:
             logger.info("Empty defaultssheet provided")

src/semeio/forward_models/design_kw/design_kw.py

@@ -1,7 +1,9 @@
 import logging
 import re
 import shlex
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping, Sequence
+from re import Match
+from typing import Any
 
 from ert import ForwardModelStepWarning
 
@@ -79,28 +81,28 @@ def find_matching_errors(
     return errors
 
 
-def is_perl(file_name, template):
+def is_perl(file_name: str, template: Sequence[str]) -> bool:
     return file_name.endswith(".pl") or template[0].find("perl") != -1
 
 
-def is_xml(file_name: str, template: list[str]) -> bool:
+def is_xml(file_name: str, template: Sequence[str]) -> bool:
     return file_name.endswith(".xml") or template[0].find("?xml") != -1
 
 
-def unmatched_templates(line):
+def unmatched_templates(line: str) -> list[str]:
     bracketpattern = re.compile("<.+?>")
     if bracketpattern.search(line):
         return bracketpattern.findall(line)
     return []
 
 
-def is_comment(line):
+def is_comment(line: str) -> Match[str] | None:
     ecl_comment_pattern = re.compile("^--")
     std_comment_pattern = re.compile("^#")
     return ecl_comment_pattern.search(line) or std_comment_pattern.search(line)
 
 
-def extract_key_value(parameters: list[str]) -> dict[str, str]:
+def extract_key_value(parameters: Iterable[str]) -> dict[str, Any]:
     """Parses a list of strings, looking for key-value pairs pr. line
     separated by whitespace, into a dictionary.
 
@@ -142,7 +144,9 @@ def extract_key_value(parameters: list[str]) -> dict[str, str]:
     return res
 
 
-def rm_genkw_prefix(paramsdict, ignoreprefixes="LOG10_"):
+def rm_genkw_prefix(
+    paramsdict: Mapping[str, Any], ignoreprefixes: str | Iterable[str] | None = "LOG10_"
+) -> dict[str, Any]:
     """Strip prefixes from keys in a dictionary.
 
     Prefix is any string before a colon. No colon means no prefix.