diff --git a/src/fmu/sumo/explorer/cache.py b/src/fmu/sumo/explorer/cache.py index 0204b63f..f76eacd8 100644 --- a/src/fmu/sumo/explorer/cache.py +++ b/src/fmu/sumo/explorer/cache.py @@ -33,3 +33,8 @@ def put(self, key, value): def has(self, key): return key in self.cache + + def clear(self): + with self.lock: + self.cache.clear() + self.access.clear() diff --git a/src/fmu/sumo/explorer/objects/__init__.py b/src/fmu/sumo/explorer/objects/__init__.py index 42c32626..33cd6861 100644 --- a/src/fmu/sumo/explorer/objects/__init__.py +++ b/src/fmu/sumo/explorer/objects/__init__.py @@ -10,6 +10,8 @@ from fmu.sumo.explorer.objects.surface import Surface from fmu.sumo.explorer.objects.polygons import Polygons from fmu.sumo.explorer.objects.table import Table +from fmu.sumo.explorer.objects.cpgrid import CPGrid +from fmu.sumo.explorer.objects.cpgrid_property import CPGridProperty from fmu.sumo.explorer.objects.iteration import Iteration from fmu.sumo.explorer.objects.iterations import Iterations from fmu.sumo.explorer.objects.realization import Realization diff --git a/src/fmu/sumo/explorer/objects/_search_context.py b/src/fmu/sumo/explorer/objects/_search_context.py index fb749263..7f96e282 100644 --- a/src/fmu/sumo/explorer/objects/_search_context.py +++ b/src/fmu/sumo/explorer/objects/_search_context.py @@ -10,14 +10,6 @@ from fmu.sumo.explorer.cache import LRUCache -_CASE_FIELDS = {"include": [], "exclude": []} - -_CHILD_FIELDS = { - "include": [], - "exclude": ["data.spec.columns", "fmu.realization.parameters"], -} - - def _gen_filter_none(): def _fn(value): return None, None @@ -289,6 +281,9 @@ def __init__( self._hits = None self._cache = LRUCache(capacity=200) self._length = None + self._select = { + "excludes": ["fmu.realization.parameters"], + } return @property @@ -324,10 +319,12 @@ def _to_sumo(self, obj, blob=None): "polygons": objects.Polygons, "surface": objects.Surface, "table": objects.Table, + "cpgrid": objects.CPGrid, + "cpgrid_property": objects.CPGridProperty }.get(cls) if constructor is None: - warnings.warn(f"No constructor for class {cls}") - constructor = objects.Child + warnings.warn(f"No constructor for class {cls}") + constructor = objects.Child return constructor(self._sumo, obj, blob) def __len__(self): @@ -490,7 +487,44 @@ async def getitem_async(self, index): uuid = self._hits[index] return await self.get_object_async(uuid) - def get_object(self, uuid: str, select: List[str] = None) -> Dict: + def select(self, sel): + """Specify what should be returned from elasticsearch. + Has the side effect of clearing the lru cache. + sel is either a single string value, a list of string value, + or a dictionary with keys "includes" and/or "excludes" and + the values are lists of strings. The string values are nested + property names. + + Args: + sel (str | List(str) | Dict(str, List[str]): select specification + Returns: + None + """ + + required = set(["class"]) + def extreq(lst): + if isinstance(lst, str): + lst = [lst] + return list(set(lst) | required) + if isinstance(sel, str): + self._select = extreq([sel]) + elif isinstance(sel, list): + self._select = extreq(sel) + elif isinstance(sel, dict): + inc = sel.get("includes") + exc = sel.get("excludes") + slct = {} + if inc is not None: + slct["includes"] = extreq(inc) + pass + if exc is not None: + slct["excludes"] = exc + pass + self._select = slct + pass + self._cache.clear() + + def get_object(self, uuid: str) -> Dict: """Get metadata object by uuid Args: @@ -505,11 +539,9 @@ def get_object(self, uuid: str, select: List[str] = None) -> Dict: query = { "query": {"ids": {"values": [uuid]}}, "size": 1, + "_source": self._select } - if select is not None: - query["_source"] = select - res = self._sumo.post("/search", json=query) hits = res.json()["hits"]["hits"] @@ -520,9 +552,7 @@ def get_object(self, uuid: str, select: List[str] = None) -> Dict: return self._to_sumo(obj) - async def get_object_async( - self, uuid: str, select: List[str] = None - ) -> Dict: + async def get_object_async(self, uuid: str) -> Dict: """Get metadata object by uuid Args: @@ -538,11 +568,9 @@ async def get_object_async( query = { "query": {"ids": {"values": [uuid]}}, "size": 1, + "_source": self._select } - if select is not None: - query["_source"] = select - res = await self._sumo.post_async("/search", json=query) hits = res.json()["hits"]["hits"] @@ -557,13 +585,11 @@ def _maybe_prefetch(self, index): uuid = self._hits[index] if self._cache.has(uuid): return - uuids = self._hits[index : min(index + 100, len(self._hits))] + uuids = self._hits[index:min(index + 100, len(self._hits))] uuids = [uuid for uuid in uuids if not self._cache.has(uuid)] hits = self.__search_all( {"ids": {"values": uuids}}, - select={ - "excludes": ["fmu.realization.parameters"], - }, + select=self._select, ) if len(hits) == 0: return @@ -575,13 +601,11 @@ async def _maybe_prefetch_async(self, index): uuid = self._hits[index] if self._cache.has(uuid): return - uuids = self._hits[index : min(index + 100, len(self._hits))] + uuids = self._hits[index:min(index + 100, len(self._hits))] uuids = [uuid for uuid in uuids if not self._cache.has(uuid)] hits = await self.__search_all_async( {"ids": {"values": uuids}}, - select={ - "excludes": ["fmu.realization.parameters"], - }, + select=self._select, ) if len(hits) == 0: return @@ -809,17 +833,38 @@ def all(self): @property def cases(self): """Cases from current selection.""" - return objects.Cases(self) + uuids = self._get_field_values("fmu.case.uuid.keyword") + return objects.Cases(self, uuids) + + @property + async def cases_async(self): + """Cases from current selection.""" + uuids = await self._get_field_values_async("fmu.case.uuid.keyword") + return objects.Cases(self, uuids) @property def iterations(self): """Iterations from current selection.""" - return objects.Iterations(self) + uuids = self._get_field_values("fmu.iteration.uuid.keyword") + return objects.Iterations(self, uuids) + + @property + async def iterations_async(self): + """Iterations from current selection.""" + uuids = await self._get_field_values_async("fmu.iteration.uuid.keyword") + return objects.Iterations(self, uuids) @property def realizations(self): """Realizations from current selection.""" - return objects.Realizations(self) + uuids = self._get_field_values("fmu.realization.uuid.keyword") + return objects.Realizations(self, uuids) + + @property + async def realizations_async(self): + """Realizations from current selection.""" + uuids = await self._get_field_values_async("fmu.realization.uuid.keyword") + return objects.Realizations(self, uuids) @property def template_paths(sc): @@ -958,6 +1003,14 @@ def polygons(self): def dictionaries(self): return self._context_for_class("dictionary") + @property + def grids(self): + return self._context_for_class("cpgrid") + + @property + def grid_properties(self): + return self._context_for_class("cpgrid_property") + def _get_object_by_class_and_uuid(self, cls, uuid): obj = self.get_object(uuid) if obj.metadata["class"] != cls: @@ -1100,7 +1153,7 @@ def get_surface_by_uuid(self, uuid: str): Returns: Surface: surface object """ - metadata = self.get_object(uuid, _CHILD_FIELDS) + metadata = self.get_object(uuid) return objects.Surface(self._sumo, metadata) async def get_surface_by_uuid_async(self, uuid: str): @@ -1112,7 +1165,7 @@ async def get_surface_by_uuid_async(self, uuid: str): Returns: Surface: surface object """ - metadata = await self.get_object_async(uuid, _CHILD_FIELDS) + metadata = await self.get_object_async(uuid) return objects.Surface(self._sumo, metadata) def get_polygons_by_uuid(self, uuid: str): @@ -1124,7 +1177,7 @@ def get_polygons_by_uuid(self, uuid: str): Returns: Polygons: polygons object """ - metadata = self.get_object(uuid, _CHILD_FIELDS) + metadata = self.get_object(uuid) return objects.Polygons(self._sumo, metadata) async def get_polygons_by_uuid_async(self, uuid: str): @@ -1136,7 +1189,7 @@ async def get_polygons_by_uuid_async(self, uuid: str): Returns: Polygons: polygons object """ - metadata = await self.get_object_async(uuid, _CHILD_FIELDS) + metadata = await self.get_object_async(uuid) return objects.Polygons(self._sumo, metadata) def get_table_by_uuid(self, uuid: str): @@ -1148,7 +1201,7 @@ def get_table_by_uuid(self, uuid: str): Returns: Table: table object """ - metadata = self.get_object(uuid, _CHILD_FIELDS) + metadata = self.get_object(uuid) return objects.Table(self._sumo, metadata) async def get_table_by_uuid_async(self, uuid: str): @@ -1160,7 +1213,7 @@ async def get_table_by_uuid_async(self, uuid: str): Returns: Table: table object """ - metadata = await self.get_object_async(uuid, _CHILD_FIELDS) + metadata = await self.get_object_async(uuid) return objects.Table(self._sumo, metadata) def _verify_aggregation_operation(self): diff --git a/src/fmu/sumo/explorer/objects/cases.py b/src/fmu/sumo/explorer/objects/cases.py index e607aec5..fa50b2e5 100644 --- a/src/fmu/sumo/explorer/objects/cases.py +++ b/src/fmu/sumo/explorer/objects/cases.py @@ -3,45 +3,19 @@ from fmu.sumo.explorer.objects._search_context import SearchContext class Cases(SearchContext): - def __init__(self, sc): - super().__init__(sc._sumo, sc._must, sc._must_not) + def __init__(self, sc, uuids): + super().__init__(sc._sumo, must=[{"terms": {"fmu.case.uuid.keyword": uuids}}]) + self._hits = uuids return - def __len__(self): - if self._length is None: - if self._hits is None: - self._hits = self._search_all(select=False) - pass - self._length = len(self._hits) - pass - return self._length - - async def length_async(self): - if self._length is None: - if self._hits is None: - self._hits = self._search_all(select=False) - pass - self._length = len(self._hits) - pass - return self._length - - def _search_all(self, select=False): - uuids = self._get_field_values("fmu.case.uuid.keyword") - if select is False: - return uuids - # ELSE - return SearchContext(must=[{"ids": {"values": uuids}}])._search_all(select=select) - - async def _search_all_async(self, select=False): - uuids = await self._get_field_values_async("fmu.case.uuid.keyword") - if select is False: - return uuids - # ELSE - return await SearchContext(must=[{"ids": {"values": uuids}}])._search_all_async(select=select) - def _maybe_prefetch(self, index): return async def _maybe_prefetch_async(self, index): return + + def filter(self, **kwargs): + sc = super().filter(**kwargs) + uuids = sc.uuids + return Cases(self, uuids) diff --git a/src/fmu/sumo/explorer/objects/cpgrid.py b/src/fmu/sumo/explorer/objects/cpgrid.py new file mode 100644 index 00000000..fceb291f --- /dev/null +++ b/src/fmu/sumo/explorer/objects/cpgrid.py @@ -0,0 +1,46 @@ +"""Module containing class for cpgrid""" + +from typing import Dict +from sumo.wrapper import SumoClient +from fmu.sumo.explorer.objects._child import Child + +class CPGrid(Child): + """Class representing a cpgrid object in Sumo.""" + + def __init__(self, sumo: SumoClient, metadata: Dict, blob=None) -> None: + """ + Args: + sumo (SumoClient): connection to Sumo + metadata (dict): dictionary metadata + blob: data object + """ + super().__init__(sumo, metadata, blob) + + def to_cpgrid(self): + """Get cpgrid object as a Grid + Returns: + Grid: A Grid object + """ + try: + from xtgeo import grid_from_file + except ModuleNotFoundError: + raise RuntimeError("Unable to import xtgeo; probably not installed.") + try: + return grid_from_file(self.blob) + except TypeError as type_err: + raise TypeError(f"Unknown format: {self.format}") from type_err + + async def to_cpgrid_async(self): + """Get cpgrid object as a Grid + Returns: + Grid: A Grid object + """ + try: + from xtgeo import grid_from_file + except ModuleNotFoundError: + raise RuntimeError("Unable to import xtgeo; probably not installed.") + + try: + return grid_from_file(await self.blob_async) + except TypeError as type_err: + raise TypeError(f"Unknown format: {self.format}") from type_err diff --git a/src/fmu/sumo/explorer/objects/cpgrid_property.py b/src/fmu/sumo/explorer/objects/cpgrid_property.py new file mode 100644 index 00000000..b146aa94 --- /dev/null +++ b/src/fmu/sumo/explorer/objects/cpgrid_property.py @@ -0,0 +1,46 @@ +"""Module containing class for cpgrid_property""" + +from typing import Dict +from sumo.wrapper import SumoClient +from fmu.sumo.explorer.objects._child import Child + +class CPGridProperty(Child): + """Class representing a cpgrid_property object in Sumo.""" + + def __init__(self, sumo: SumoClient, metadata: Dict, blob=None) -> None: + """ + Args: + sumo (SumoClient): connection to Sumo + metadata (dict): dictionary metadata + blob: data object + """ + super().__init__(sumo, metadata, blob) + + def to_cpgrid_property(self): + """Get cpgrid_property object as a GridProperty + Returns: + GridProperty: A GridProperty object + """ + try: + from xtgeo import gridproperty_from_file + except ModuleNotFoundError: + raise RuntimeError("Unable to import xtgeo; probably not installed.") + try: + return gridproperty_from_file(self.blob) + except TypeError as type_err: + raise TypeError(f"Unknown format: {self.format}") from type_err + + async def to_cpgrid_property_async(self): + """Get cpgrid_property object as a GridProperty + Returns: + GridProperty: A GridProperty object + """ + try: + from xtgeo import gridproperty_from_file + except ModuleNotFoundError: + raise RuntimeError("Unable to import xtgeo; probably not installed.") + + try: + return gridproperty_from_file(await self.blob_async) + except TypeError as type_err: + raise TypeError(f"Unknown format: {self.format}") from type_err diff --git a/src/fmu/sumo/explorer/objects/iterations.py b/src/fmu/sumo/explorer/objects/iterations.py index fecb0d9f..bdee9306 100644 --- a/src/fmu/sumo/explorer/objects/iterations.py +++ b/src/fmu/sumo/explorer/objects/iterations.py @@ -4,34 +4,11 @@ from fmu.sumo.explorer.objects._search_context import SearchContext class Iterations(SearchContext): - def __init__(self, _search_context): - super().__init__(_search_context._sumo, _search_context._must, _search_context._must_not) + def __init__(self, sc, uuids): + super().__init__(sc._sumo, must=[{"terms": {"fmu.iteration.uuid.keyword": uuids}}]) + self._hits = uuids return - def __len__(self): - if self._length is None: - if self._hits is None: - self._hits = self._search_all(select=False) - pass - self._length = len(self._hits) - pass - return self._length - - async def length_async(self): - if self._length is None: - if self._hits is None: - self._hits = self._search_all(select=False) - pass - self._length = len(self._hits) - pass - return self._length - - def _search_all(self, select=False): - return self._get_field_values("fmu.iteration.uuid.keyword") - - async def _search_all_async(self, select=False): - return await self._get_field_values_async("fmu.iteration.uuid.keyword") - def _maybe_prefetch(self, index): return @@ -76,3 +53,7 @@ async def get_object_async( return obj + def filter(self, **kwargs): + sc = super().filter(**kwargs) + uuids = sc.uuids + return Iterations(self, uuids) diff --git a/src/fmu/sumo/explorer/objects/realizations.py b/src/fmu/sumo/explorer/objects/realizations.py index b328cddb..405ab644 100644 --- a/src/fmu/sumo/explorer/objects/realizations.py +++ b/src/fmu/sumo/explorer/objects/realizations.py @@ -4,33 +4,11 @@ from fmu.sumo.explorer.objects._search_context import SearchContext class Realizations(SearchContext): - def __init__(self, sc): - super().__init__(sc._sumo, sc._must, sc._must_not) + def __init__(self, sc, uuids): + super().__init__(sc._sumo, must=[{"terms": {"fmu.realization.uuid.keyword": uuids}}]) + self._hits = uuids return - def __len__(self): - if self._length is None: - if self._hits is None: - self._hits = self._search_all(select=False) - pass - self._length = len(self._hits) - pass - return self._length - - async def length_async(self): - if self._length is None: - if self._hits is None: - self._hits = self._search_all(select=False) - pass - self._length = len(self._hits) - pass - return self._length - - def _search_all(self, select=False): - return self._get_field_values("fmu.realization.uuid.keyword") - - async def _search_all_async(self, select=False): - return await self._get_field_values_async("fmu.realization.uuid.keyword") def _maybe_prefetch(self, index): return @@ -75,3 +53,7 @@ async def get_object_async( return obj + def filter(self, **kwargs): + sc = super().filter(**kwargs) + uuids = sc.uuids + return Realizations(self, uuids) diff --git a/tests/test_explorer.py b/tests/test_explorer.py index 9b9fda67..293d6f3d 100644 --- a/tests/test_explorer.py +++ b/tests/test_explorer.py @@ -327,4 +327,39 @@ def test_seismic_case_by_uuid(explorer: Explorer, seismic_case_uuid: str): assert "Trace" in channel_list assert "SEGYTraceHeader" in channel_list - +def test_grids_and_properties(explorer: Explorer): + cases_with_grids = explorer.grids.cases.filter(status="keep") + cases_with_gridprops = explorer.grid_properties.cases.filter(status="keep") + cgs=set([case.uuid for case in cases_with_grids]) + cgps=set([case.uuid for case in cases_with_gridprops]) + assert cgs==cgps + case=cases_with_grids[0] + grids=case.grids + gridprops=case.grid_properties + xtgrid=grids[0].to_cpgrid() + gridspec=grids[0].metadata["data"]["spec"] + assert xtgrid.nlay == gridspec["nlay"] + assert xtgrid.nrow == gridspec["nrow"] + assert xtgrid.ncol == gridspec["ncol"] + xtgridprop=gridprops[0].to_cpgrid_property() + gridpropspec = gridprops[0].metadata["data"]["spec"] + assert xtgridprop.nlay == gridpropspec["nlay"] + assert xtgridprop.nrow == gridpropspec["nrow"] + assert xtgridprop.ncol == gridpropspec["ncol"] + +def test_search_context_select(test_case: Case): + surfs = test_case.surfaces.filter(realization=True) + assert "_sumo" in surfs[0].metadata + surfs.select("fmu") + assert "_sumo" not in surfs[0].metadata + assert "fmu" in surfs[0].metadata + surfs.select(["fmu"]) + assert "_sumo" not in surfs[0].metadata + assert "fmu" in surfs[0].metadata + surfs.select({"excludes": ["fmu"]}) + assert "_sumo" in surfs[0].metadata + assert "fmu" not in surfs[0].metadata + surfs.select({"includes": ["_sumo"], "excludes": ["_sumo.timestamp"]}) + assert "_sumo" in surfs[0].metadata + assert "fmu" not in surfs[0].metadata + assert "timestamp" not in surfs[0].metadata["_sumo"]