From 86499ad384bb20d3f6d0409d672ea996f598e829 Mon Sep 17 00:00:00 2001 From: JB Lovland Date: Fri, 1 Mar 2024 10:31:01 +0100 Subject: [PATCH] CLN: Remove last usages of the Any type --- src/fmu/dataio/_filedata_provider.py | 19 ++++++++--- src/fmu/dataio/_metadata.py | 19 +++++++---- src/fmu/dataio/_objectdata_provider.py | 46 +++++++++++++++----------- src/fmu/dataio/dataio.py | 6 ++-- 4 files changed, 57 insertions(+), 33 deletions(-) diff --git a/src/fmu/dataio/_filedata_provider.py b/src/fmu/dataio/_filedata_provider.py index 6b79cd326..feb38cb69 100644 --- a/src/fmu/dataio/_filedata_provider.py +++ b/src/fmu/dataio/_filedata_provider.py @@ -9,7 +9,7 @@ from copy import deepcopy from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Optional from warnings import warn from ._definitions import FmuContext @@ -17,6 +17,10 @@ logger: Final = null_logger(__name__) +if TYPE_CHECKING: + from ._objectdata_provider import ObjectDataProvider + from .dataio import ExportData + @dataclass class FileDataProvider: @@ -31,8 +35,8 @@ class FileDataProvider: """ # input - dataio: Any - objdata: Any + dataio: ExportData + objdata: ObjectDataProvider rootpath: Path = field(default_factory=Path) itername: str = "" realname: str = "" @@ -161,8 +165,13 @@ def _get_path(self) -> tuple[Path, Path | None]: """Construct and get the folder path(s).""" linkdest = None + assert isinstance( + self.dataio.fmu_context, FmuContext + ) # Converted to a FmuContext obj. in post-init. + dest = self._get_path_generic( - mode=self.dataio.fmu_context, allow_forcefolder=True + mode=self.dataio.fmu_context, + allow_forcefolder=True, ) if self.dataio.fmu_context == FmuContext.CASE_SYMLINK_REALIZATION: @@ -176,7 +185,7 @@ def _get_path(self) -> tuple[Path, Path | None]: def _get_path_generic( self, - mode: Literal[FmuContext.REALIZATION, FmuContext.PREPROCESSED], + mode: FmuContext, allow_forcefolder: bool = True, info: str = "", ) -> Path: diff --git a/src/fmu/dataio/_metadata.py b/src/fmu/dataio/_metadata.py index 15d927ec7..3083eccee 100644 --- a/src/fmu/dataio/_metadata.py +++ b/src/fmu/dataio/_metadata.py @@ -14,7 +14,7 @@ from datetime import timezone from pathlib import Path from tempfile import NamedTemporaryFile -from typing import Any, Final +from typing import Final from warnings import warn from fmu import dataio @@ -30,6 +30,7 @@ ) from fmu.dataio.datastructure.meta import meta +from . import types from ._definitions import FmuContext from ._logging import null_logger @@ -220,13 +221,13 @@ class MetaData: """ # input variables - obj: Any - dataio: Any + obj: types.Inferrable + dataio: dataio.ExportData compute_md5: bool = True # storage state variables - objdata: Any = field(default=None, init=False) - fmudata: Any = field(default=None, init=False) + objdata: ObjectDataProvider | None = field(default=None, init=False) + fmudata: FmuProvider | None = field(default=None, init=False) iter_name: str = field(default="", init=False) real_name: str = field(default="", init=False) @@ -281,8 +282,8 @@ def _populate_meta_fmu(self) -> None: """ fmudata = FmuProvider( model=self.dataio.config.get("model", None), - fmu_context=self.dataio.fmu_context, - casepath_proposed=self.dataio.casepath, + fmu_context=FmuContext.get(self.dataio.fmu_context), + casepath_proposed=self.dataio.casepath or "", include_ertjobs=self.dataio.include_ertjobs, forced_realization=self.dataio.realization, workflow=self.dataio.workflow, @@ -328,6 +329,8 @@ def _populate_meta_file(self) -> None: - absolute_path_symlink, as above but full path """ + assert self.objdata is not None + fdata = FileDataProvider( self.dataio, self.objdata, @@ -362,6 +365,7 @@ def _populate_meta_file(self) -> None: def _populate_meta_class(self) -> None: """Get the general class which is a simple string.""" + assert self.objdata is not None self.meta_class = self.objdata.classname def _populate_meta_tracklog(self) -> None: @@ -395,6 +399,7 @@ def _populate_meta_display(self) -> None: if self.dataio.display_name is not None: display_name = self.dataio.display_name else: + assert self.objdata is not None display_name = self.objdata.name self.meta_display = {"name": display_name} diff --git a/src/fmu/dataio/_objectdata_provider.py b/src/fmu/dataio/_objectdata_provider.py index bbd530806..01e81403e 100644 --- a/src/fmu/dataio/_objectdata_provider.py +++ b/src/fmu/dataio/_objectdata_provider.py @@ -147,15 +147,18 @@ class DerivedObjectDescriptor: "table", "dictionary", ] - efolder: Literal[ - "maps", - "polygons", - "points", - "cubes", - "grids", - "tables", - "dictionaries", - ] + efolder: ( + Literal[ + "maps", + "polygons", + "points", + "cubes", + "grids", + "tables", + "dictionaries", + ] + | str + ) fmt: str extension: str spec: Dict[str, Any] @@ -230,8 +233,8 @@ class ObjectDataProvider: """ # input fields - obj: Any - dataio: Any + obj: types.Inferrable + dataio: dataio.ExportData meta_existing: dict = field(default_factory=dict) # result properties; the most important is metadata which IS the 'data' part in @@ -262,7 +265,7 @@ def _derive_name_stratigraphy(self) -> DerivedNamedStratigraphy: name = derive_name(self.dataio, self.obj) # next check if usename has a "truename" and/or aliases from the config - strat = self.dataio.config.get("stratigraphy") # shortform + strat = self.dataio.config.get("stratigraphy", {}) no_start_or_missing_name = strat is None or name not in strat rv = DerivedNamedStratigraphy( @@ -658,11 +661,11 @@ def _derive_spec_bbox_dataframe( ) -> SpecificationAndBoundingBox: """Process/collect the data items for DataFrame.""" logger.info("Process data metadata for DataFrame (tables)") - df: pd.DataFrame = self.obj + assert isinstance(self.obj, pd.DataFrame) return SpecificationAndBoundingBox( spec=specification.TableSpecification( - columns=list(df.columns), - size=int(df.size), + columns=list(self.obj.columns), + size=int(self.obj.size), ).model_dump( mode="json", exclude_none=True, @@ -675,11 +678,13 @@ def _derive_spec_bbox_arrowtable( ) -> SpecificationAndBoundingBox: """Process/collect the data items for Arrow table.""" logger.info("Process data metadata for arrow (tables)") - table = self.obj + from pyarrow import Table + + assert isinstance(self.obj, Table) return SpecificationAndBoundingBox( spec=specification.TableSpecification( - columns=list(table.column_names), - size=table.num_columns * table.num_rows, + columns=list(self.obj.column_names), + size=self.obj.num_columns * self.obj.num_rows, ).model_dump( mode="json", exclude_none=True, @@ -699,6 +704,9 @@ def _get_columns(self) -> list[str]: columns = list(self.obj.columns) else: logger.debug("arrow") + from pyarrow import Table + + assert isinstance(self.obj, Table) columns = self.obj.column_names logger.debug("Available columns in table %s ", columns) return columns @@ -813,7 +821,7 @@ def _derive_from_existing(self) -> None: # TODO: Clean up types below. self.time0, self.time1 = parse_timedata(self.meta_existing["data"]) # type: ignore - def _process_content(self) -> tuple[str, dict | None]: + def _process_content(self) -> tuple[str | dict, dict | None]: """Work with the `content` metadata""" # content == "unset" is not wanted, but in case metadata has been produced while diff --git a/src/fmu/dataio/dataio.py b/src/fmu/dataio/dataio.py index 8c2ebe211..60f669e2a 100644 --- a/src/fmu/dataio/dataio.py +++ b/src/fmu/dataio/dataio.py @@ -426,7 +426,9 @@ class ExportData: depth_reference: str = "msl" description: Union[str, list] = "" display_name: Optional[str] = None - fmu_context: Union[FmuContext, str] = "realization" + fmu_context: Union[FmuContext, str] = ( + FmuContext.REALIZATION + ) # post init converts to FmuContext forcefolder: str = "" grid_model: Optional[str] = None is_observation: bool = False @@ -447,7 +449,7 @@ class ExportData: table_index: Optional[list] = None # some keys that are modified version of input, prepended with _use - _usecontent: dict = field(default_factory=dict, init=False) + _usecontent: dict | str = field(default_factory=dict, init=False) _usefmtflag: str = field(default="", init=False) # storing resulting state variables for instance, non-public: