Skip to content

Commit

Permalink
CLN: Add typing to _grd3d_utils
Browse files Browse the repository at this point in the history
  • Loading branch information
mferrera committed Nov 1, 2023
1 parent 3972509 commit 734b173
Showing 1 changed file with 31 additions and 19 deletions.
50 changes: 31 additions & 19 deletions src/xtgeo/grid3d/_grid3d_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
# -*- coding: utf-8 -*-

"""Some grid utilities, file scanning etc (methods with no class)"""

from __future__ import annotations

import re
from typing import TYPE_CHECKING, List, Literal, Union

import pandas as pd

Expand All @@ -10,13 +12,22 @@
from xtgeo import XTGeoCLibError
from xtgeo.common.constants import MAXDATES, MAXKEYWORDS

if TYPE_CHECKING:
from xtgeo.common.sys import _XTGeoFile

from .grid_properties import GridPropertiesKeywords, KeywordDateTuple, KeywordTuple

xtg = xtgeo.XTGeoDialog()
logger = xtg.functionlogger(__name__)


def scan_keywords(
pfile, fformat="xecl", maxkeys=MAXKEYWORDS, dataframe=False, dates=False
):
pfile: _XTGeoFile,
fformat: Literal["xecl", "roff"] = "xecl",
maxkeys: int = MAXKEYWORDS,
dataframe: bool = False,
dates: bool = False,
) -> GridPropertiesKeywords:
"""Quick scan of keywords in Eclipse binary restart/init/... file,
or ROFF binary files.
Expand All @@ -27,21 +38,20 @@ def scan_keywords(

if fformat == "xecl":
if dates:
data = _scan_ecl_keywords_w_dates(
keywords = _scan_ecl_keywords_w_dates(
pfile, maxkeys=maxkeys, dataframe=dataframe
)
else:
data = _scan_ecl_keywords(pfile, maxkeys=maxkeys, dataframe=dataframe)

keywords = _scan_ecl_keywords(pfile, maxkeys=maxkeys, dataframe=dataframe)
elif fformat == "roff":
keywords = _scan_roff_keywords(pfile, maxkeys=maxkeys, dataframe=dataframe)
else:
data = _scan_roff_keywords(pfile, maxkeys=maxkeys, dataframe=dataframe)

raise ValueError(f"File format can be either `roff` or `xecl`, given {fformat}")
pfile.cfclose()

return data
return keywords


def scan_dates(pfile, maxdates=MAXDATES, dataframe=False):
def scan_dates(pfile: _XTGeoFile, maxdates: int = MAXDATES, dataframe: bool = False):
"""Quick scan dates in a simulation restart file.
Cf. grid_properties.py description
Expand Down Expand Up @@ -81,8 +91,9 @@ def scan_dates(pfile, maxdates=MAXDATES, dataframe=False):
return zdates


def _scan_ecl_keywords(pfile, maxkeys=MAXKEYWORDS, dataframe=False):
logger.info("Scanning ECL keywords...")
def _scan_ecl_keywords(
pfile: _XTGeoFile, maxkeys: int = MAXKEYWORDS, dataframe: bool = False
) -> Union[List[KeywordTuple], pd.DataFrame]:
cfhandle = pfile.get_cfhandle()

# maxkeys*10 is used for 1D kewords; 10 => max 8 letters in eclipse +
Expand Down Expand Up @@ -131,12 +142,11 @@ def _scan_ecl_keywords(pfile, maxkeys=MAXKEYWORDS, dataframe=False):
return result


def _scan_ecl_keywords_w_dates(pfile, maxkeys=MAXKEYWORDS, dataframe=False):
def _scan_ecl_keywords_w_dates(
pfile: _XTGeoFile, maxkeys: int = MAXKEYWORDS, dataframe: bool = False
) -> Union[List[KeywordDateTuple], pd.DataFrame]:
"""Add a date column to the keyword"""

logger.info("Scan keywords with dates...")
xkeys = _scan_ecl_keywords(pfile, maxkeys=maxkeys, dataframe=False)

xdates = scan_dates(pfile, maxdates=MAXDATES, dataframe=False)

result = []
Expand All @@ -160,7 +170,9 @@ def _scan_ecl_keywords_w_dates(pfile, maxkeys=MAXKEYWORDS, dataframe=False):
return result


def _scan_roff_keywords(pfile, maxkeys=MAXKEYWORDS, dataframe=False):
def _scan_roff_keywords(
pfile: _XTGeoFile, maxkeys: int = MAXKEYWORDS, dataframe: bool = False
) -> Union[List[KeywordTuple], pd.DataFrame]:
rectypes = _cxtgeo.new_intarray(maxkeys)
reclens = _cxtgeo.new_longarray(maxkeys)
recstarts = _cxtgeo.new_longarray(maxkeys)
Expand Down

0 comments on commit 734b173

Please sign in to comment.