fmu-sumo 2.3.6__py3-none-any.whl → 2.3.8__py3-none-any.whl

fmu/sumo/explorer/objects/_search_context.py

@@ -1,19 +1,27 @@
+from __future__ import annotations
+
 import uuid
 import warnings
 from datetime import datetime
 from io import BytesIO
-from typing import Dict, List, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
 
 import deprecation
 import httpx
-from sumo.wrapper import SumoClient
 
-import fmu.sumo.explorer.objects as objects
+from fmu.sumo.explorer import objects
 from fmu.sumo.explorer.cache import LRUCache
 
+if TYPE_CHECKING:
+    from sumo.wrapper import SumoClient
+
+
+# Type aliases
+SelectArg = Union[bool, str, Dict[str, Union[str, List[str]]], List[str]]
+
 
 def _gen_filter_none():
-    def _fn(value):
+    def _fn(_):
         return None, None
 
     return _fn
@@ -211,7 +219,7 @@ class Pit:
         self._id = res.json()["id"]
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, *_):
         if self._id is not None:
             self._sumo.delete("/pit", params={"id": self._id})
             pass
@@ -224,7 +232,7 @@ class Pit:
         self._id = res.json()["id"]
         return self
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
+    async def __aexit__(self, *_):
         if self._id is not None:
             await self._sumo.delete_async("/pit", params={"id": self._id})
             pass
@@ -257,7 +265,7 @@ class SearchContext:
         self._hits = None
         self._cache = LRUCache(capacity=200)
         self._length = None
-        self._select = {
+        self._select: SelectArg = {
             "excludes": ["fmu.realization.parameters"],
         }
         return
@@ -282,7 +290,7 @@ class SearchContext:
             else:
                 return {"bool": {"must": must, "must_not": must_not}}
 
-    def _to_sumo(self, obj, blob=None):
+    def _to_sumo(self, obj, blob=None) -> objects.Document:
         cls = obj["_source"]["class"]
         if cls == "case":
             return objects.Case(self._sumo, obj)
@@ -295,6 +303,8 @@ class SearchContext:
             "table": objects.Table,
             "cpgrid": objects.CPGrid,
             "cpgrid_property": objects.CPGridProperty,
+            "iteration": objects.Iteration,
+            "realization": objects.Realization,
         }.get(cls)
         if constructor is None:
             warnings.warn(f"No constructor for class {cls}")
@@ -319,7 +329,7 @@ class SearchContext:
             self._length = res["count"]
         return self._length
 
-    def __search_all(self, query, size=1000, select=False):
+    def __search_all(self, query, size: int = 1000, select: SelectArg = False):
         all_hits = []
         query = {
             "query": query,
@@ -356,10 +366,12 @@ class SearchContext:
             pass
         return all_hits
 
-    def _search_all(self, select=False):
+    def _search_all(self, select: SelectArg = False):
         return self.__search_all(query=self._query, size=1000, select=select)
 
-    async def __search_all_async(self, query, size=1000, select=False):
+    async def __search_all_async(
+        self, query, size: int = 1000, select: SelectArg = False
+    ):
         all_hits = []
         query = {
             "query": query,
@@ -398,7 +410,7 @@ class SearchContext:
             pass
         return all_hits
 
-    async def _search_all_async(self, select=False):
+    async def _search_all_async(self, select: SelectArg = False):
         return await self.__search_all_async(
             query=self._query, size=1000, select=select
         )
@@ -477,7 +489,7 @@ class SearchContext:
         assert await self.length_async() == 1
         return await self.getitem_async(0)
 
-    def select(self, sel):
+    def select(self, sel) -> SearchContext:
         """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,
@@ -521,7 +533,7 @@ class SearchContext:
         self._cache.clear()
         return self
 
-    def get_object(self, uuid: str) -> Dict:
+    def get_object(self, uuid: str) -> objects.Document:
         """Get metadata object by uuid
 
         Args:
@@ -549,7 +561,7 @@ class SearchContext:
 
         return self._to_sumo(obj)
 
-    async def get_object_async(self, uuid: str) -> Dict:
+    async def get_object_async(self, uuid: str) -> objects.Document:
         """Get metadata object by uuid
 
         Args:
@@ -579,6 +591,7 @@ class SearchContext:
         return self._to_sumo(obj)
 
     def _maybe_prefetch(self, index):
+        assert isinstance(self._hits, list)
         uuid = self._hits[index]
         if self._cache.has(uuid):
             return
@@ -595,6 +608,7 @@ class SearchContext:
         return
 
     async def _maybe_prefetch_async(self, index):
+        assert isinstance(self._hits, list)
         uuid = self._hits[index]
         if self._cache.has(uuid):
             return
@@ -613,13 +627,13 @@ class SearchContext:
     def get_objects(
         self,
         uuids: List[str],
-        select: List[str] = None,
+        select: SelectArg,
     ) -> List[Dict]:
         size = (
             1000
             if select is False
             else 100
-            if isinstance(select, list)
+            if isinstance(select, (list, dict))
             else 10
         )
         return self.__search_all(
@@ -627,15 +641,13 @@ class SearchContext:
         )
 
     async def get_objects_async(
-        self,
-        uuids: List[str],
-        select: List[str] = None,
+        self, uuids: List[str], select: SelectArg
     ) -> List[Dict]:
         size = (
             1000
             if select is False
             else 100
-            if isinstance(select, list)
+            if isinstance(select, (list, dict))
             else 10
         )
         return await self.__search_all_async(
@@ -900,8 +912,8 @@ class SearchContext:
         return objects.Realizations(self, uuids)
 
     @property
-    def template_paths(search_context):  # noqa: N805
-        return {obj.template_path for obj in search_context}
+    def template_paths(self) -> List[str]:
+        return {obj.template_path for obj in self}
 
     @property
     def metrics(self):
@@ -919,9 +931,9 @@ class SearchContext:
     @property
     async def timestamps_async(self) -> List[str]:
         """List of unique timestamps in SearchContext"""
-        ts = await self.get_field_values_async(
-            "data.time.t0.value", self._timestamp_query
-        )
+        ts = await self.filter(
+            complex=self._timestamp_query
+        ).get_field_values_async("data.time.t0.value")
         return [datetime.fromtimestamp(t / 1000).isoformat() for t in ts]
 
     def _extract_intervals(self, res):
@@ -987,7 +999,6 @@ class SearchContext:
             f = filters.get(k)
             if f is None:
                 raise Exception(f"Don't know how to generate filter for {k}")
-                pass
             _must, _must_not = f(v)
             if _must:
                 must.append(_must)
@@ -1023,46 +1034,46 @@ class SearchContext:
         return sc
 
     @property
-    def surfaces(self):
+    def surfaces(self) -> SearchContext:
         return self._context_for_class("surface")
 
     @property
-    def tables(self):
+    def tables(self) -> SearchContext:
         return self._context_for_class("table")
 
     @property
-    def cubes(self):
+    def cubes(self) -> SearchContext:
         return self._context_for_class("cube")
 
     @property
-    def polygons(self):
+    def polygons(self) -> SearchContext:
         return self._context_for_class("polygons")
 
     @property
-    def dictionaries(self):
+    def dictionaries(self) -> SearchContext:
         return self._context_for_class("dictionary")
 
     @property
-    def grids(self):
+    def grids(self) -> SearchContext:
         return self._context_for_class("cpgrid")
 
     @property
-    def grid_properties(self):
+    def grid_properties(self) -> SearchContext:
         return self._context_for_class("cpgrid_property")
 
-    def _get_object_by_class_and_uuid(self, cls, uuid):
+    def _get_object_by_class_and_uuid(self, cls, uuid) -> Any:
         obj = self.get_object(uuid)
         if obj.metadata["class"] != cls:
             raise Exception(f"Document of type {cls} not found: {uuid}")
         return obj
 
-    async def _get_object_by_class_and_uuid_async(self, cls, uuid):
+    async def _get_object_by_class_and_uuid_async(self, cls, uuid) -> Any:
         obj = await self.get_object_async(uuid)
         if obj.metadata["class"] != cls:
             raise Exception(f"Document of type {cls} not found: {uuid}")
         return obj
 
-    def get_case_by_uuid(self, uuid: str):
+    def get_case_by_uuid(self, uuid: str) -> objects.Case:
         """Get case object by uuid
 
         Args:
@@ -1073,7 +1084,7 @@ class SearchContext:
         """
         return self._get_object_by_class_and_uuid("case", uuid)
 
-    async def get_case_by_uuid_async(self, uuid: str):
+    async def get_case_by_uuid_async(self, uuid: str) -> objects.Case:
         """Get case object by uuid
 
         Args:
@@ -1091,6 +1102,7 @@ class SearchContext:
             "_source": {
                 "includes": [
                     "$schema",
+                    "class",
                     "source",
                     "version",
                     "access",
@@ -1101,7 +1113,7 @@ class SearchContext:
             },
         }
 
-    def get_iteration_by_uuid(self, uuid: str):
+    def get_iteration_by_uuid(self, uuid: str) -> objects.Iteration:
         """Get iteration object by uuid
 
         Args:
@@ -1109,14 +1121,25 @@ class SearchContext:
 
         Returns: iteration object
         """
-        res = self._sumo.post(
-            "/search", json=self._iteration_query(uuid)
-        ).json()
-        obj = res["hits"]["hits"][0]
-        obj["_id"] = uuid
-        return objects.Iteration(self._sumo, obj)
+        try:
+            obj = self.get_object(uuid)
+            assert isinstance(obj, objects.Iteration)
+            return obj
+        except Exception:
+            res = self._sumo.post(
+                "/search", json=self._iteration_query(uuid)
+            ).json()
+            hits = res["hits"]["hits"]
+            if len(hits) == 0:
+                raise Exception(f"Document not found: {uuid}")
+            obj = hits[0]
+            obj["_id"] = uuid
+            obj["_source"]["class"] = "iteration"
+            return self._to_sumo(obj)
 
-    async def get_iteration_by_uuid_async(self, uuid: str):
+    async def get_iteration_by_uuid_async(
+        self, uuid: str
+    ) -> objects.Iteration:
         """Get iteration object by uuid
 
         Args:
@@ -1124,16 +1147,25 @@ class SearchContext:
 
         Returns: iteration object
         """
-        res = (
-            await self._sumo.post_async(
-                "/search", json=self._iteration_query(uuid)
-            )
-        ).json()
-        obj = res["hits"]["hits"][0]
-        obj["_id"] = uuid
-        return objects.Iteration(self._sumo, obj)
+        try:
+            obj = await self.get_object_async(uuid)
+            assert isinstance(obj, objects.Iteration)
+            return obj
+        except Exception:
+            res = (
+                await self._sumo.post_async(
+                    "/search", json=self._iteration_query(uuid)
+                )
+            ).json()
+            hits = res["hits"]["hits"]
+            if len(hits) == 0:
+                raise Exception(f"Document not found: {uuid}")
+            obj = hits[0]
+            obj["_id"] = uuid
+            obj["_source"]["class"] = "iteration"
+            return self._to_sumo(obj)
 
-    def _realization_query(self, uuid):
+    def _realization_query(self, uuid) -> Dict:
         return {
             "query": {
                 "term": {"fmu.realization.uuid.keyword": {"value": uuid}}
@@ -1142,6 +1174,7 @@ class SearchContext:
             "_source": {
                 "includes": [
                     "$schema",
+                    "class",
                     "source",
                     "version",
                     "access",
@@ -1153,7 +1186,7 @@ class SearchContext:
             },
         }
 
-    def get_realization_by_uuid(self, uuid: str):
+    def get_realization_by_uuid(self, uuid: str) -> objects.Realization:
         """Get realization object by uuid
 
         Args:
@@ -1161,14 +1194,25 @@ class SearchContext:
 
         Returns: realization object
         """
-        res = self._sumo.post(
-            "/search", json=self._realization_query(uuid)
-        ).json()
-        obj = res["hits"]["hits"][0]
-        obj["_id"] = uuid
-        return objects.Realization(self._sumo, obj)
+        try:
+            obj = self.get_object(uuid)
+            assert isinstance(obj, objects.Realization)
+            return obj
+        except Exception:
+            res = self._sumo.post(
+                "/search", json=self._realization_query(uuid)
+            ).json()
+            hits = res["hits"]["hits"]
+            if len(hits) == 0:
+                raise Exception(f"Document not found: {uuid}")
+            obj = hits[0]
+            obj["_id"] = uuid
+            obj["_source"]["class"] = "realization"
+            return self._to_sumo(obj)
 
-    async def get_realization_by_uuid_async(self, uuid: str):
+    async def get_realization_by_uuid_async(
+        self, uuid: str
+    ) -> objects.Realization:
         """Get realization object by uuid
 
         Args:
@@ -1176,16 +1220,25 @@ class SearchContext:
 
         Returns: realization object
         """
-        res = (
-            await self._sumo.post(
-                "/search", json=self._realization_query(uuid)
-            )
-        ).json()
-        obj = res["hits"]["hits"][0]
-        obj["_id"] = uuid
-        return objects.Realization(self._sumo, obj)
+        try:
+            obj = await self.get_object_async(uuid)
+            assert isinstance(obj, objects.Realization)
+            return obj
+        except Exception:
+            res = (
+                await self._sumo.post_async(
+                    "/search", json=self._realization_query(uuid)
+                )
+            ).json()
+            hits = res["hits"]["hits"]
+            if len(hits) == 0:
+                raise Exception(f"Document not found: {uuid}")
+            obj = hits[0]
+            obj["_id"] = uuid
+            obj["_source"]["class"] = "realization"
+            return self._to_sumo(obj)
 
-    def get_surface_by_uuid(self, uuid: str):
+    def get_surface_by_uuid(self, uuid: str) -> objects.Surface:
         """Get surface object by uuid
 
         Args:
@@ -1196,7 +1249,7 @@ class SearchContext:
         """
         return self._get_object_by_class_and_uuid("surface", uuid)
 
-    async def get_surface_by_uuid_async(self, uuid: str):
+    async def get_surface_by_uuid_async(self, uuid: str) -> objects.Surface:
         """Get surface object by uuid
 
         Args:
@@ -1207,7 +1260,7 @@ class SearchContext:
         """
         return await self._get_object_by_class_and_uuid_async("surface", uuid)
 
-    def get_polygons_by_uuid(self, uuid: str):
+    def get_polygons_by_uuid(self, uuid: str) -> objects.Polygons:
         """Get polygons object by uuid
 
         Args:
@@ -1218,7 +1271,7 @@ class SearchContext:
         """
         return self._get_object_by_class_and_uuid("polygons", uuid)
 
-    async def get_polygons_by_uuid_async(self, uuid: str):
+    async def get_polygons_by_uuid_async(self, uuid: str) -> objects.Polygons:
         """Get polygons object by uuid
 
         Args:
@@ -1229,7 +1282,7 @@ class SearchContext:
         """
         return await self._get_object_by_class_and_uuid_async("polygons", uuid)
 
-    def get_table_by_uuid(self, uuid: str):
+    def get_table_by_uuid(self, uuid: str) -> objects.Table:
         """Get table object by uuid
 
         Args:
@@ -1240,7 +1293,7 @@ class SearchContext:
         """
         return self._get_object_by_class_and_uuid("table", uuid)
 
-    async def get_table_by_uuid_async(self, uuid: str):
+    async def get_table_by_uuid_async(self, uuid: str) -> objects.Table:
         """Get table object by uuid
 
         Args:
@@ -1251,7 +1304,9 @@ class SearchContext:
         """
         return await self._get_object_by_class_and_uuid_async("table", uuid)
 
-    def _verify_aggregation_operation(self):
+    def _verify_aggregation_operation(
+        self,
+    ) -> Tuple[Dict, List[str], List[int]]:
         query = {
             "query": self._query,
             "size": 1,
@@ -1269,6 +1324,8 @@ class SearchContext:
             },
         }
         sres = self._sumo.post("/search", json=query).json()
+        if len(sres["hits"]["hits"]) == 0:
+            raise Exception("No matching realizations found.")
         prototype = sres["hits"]["hits"][0]
         conflicts = [
             k
@@ -1291,8 +1348,15 @@ class SearchContext:
         rids = [hit["_source"]["fmu"]["realization"]["id"] for hit in hits]
         return prototype, uuids, rids
 
-    def _aggregate(self, columns=None, operation=None):
-        prototype, uuids, rids = self._verify_aggregation_operation()
+    def _aggregate(self, columns=None, operation=None) -> objects.Child:
+        assert (
+            operation != "collection"
+            or columns is not None
+            and len(columns) == 1
+        ), "Exactly one column required for collection aggregation."
+        prototype, uuids, rids = self.filter(
+            column=columns
+        )._verify_aggregation_operation()
         spec = {
             "object_ids": uuids,
             "operations": [operation],
@@ -1330,10 +1394,11 @@ class SearchContext:
             raise ex
         blob = BytesIO(res.content)
         res = self._to_sumo(prototype, blob)
+        assert isinstance(res, objects.Child)
         res._blob = blob
         return res
 
-    def aggregate(self, columns=None, operation=None):
+    def aggregate(self, columns=None, operation=None) -> objects.Child:
         if len(self.hidden) > 0:
             return self.hidden._aggregate(columns=columns, operation=operation)
         else:
@@ -1341,7 +1406,9 @@ class SearchContext:
                 columns=columns, operation=operation
             )
 
-    async def _verify_aggregation_operation_async(self):
+    async def _verify_aggregation_operation_async(
+        self,
+    ) -> Tuple[Dict, List[str], List[int]]:
         query = {
             "query": self._query,
             "size": 1,
@@ -1359,6 +1426,8 @@ class SearchContext:
             },
         }
         sres = (await self._sumo.post_async("/search", json=query)).json()
+        if len(sres["hits"]["hits"]) == 0:
+            raise Exception("No matching realizations found.")
         prototype = sres["hits"]["hits"][0]
         conflicts = [
             k
@@ -1381,12 +1450,21 @@ class SearchContext:
         rids = [hit["_source"]["fmu"]["realization"]["id"] for hit in hits]
         return prototype, uuids, rids
 
-    async def _aggregate_async(self, columns=None, operation=None):
+    async def _aggregate_async(
+        self, columns=None, operation=None
+    ) -> objects.Child:
+        assert (
+            operation != "collection"
+            or columns is not None
+            and len(columns) == 1
+        ), "Exactly one column required for collection aggregation."
         (
             prototype,
             uuids,
             rids,
-        ) = await self._verify_aggregation_operation_async()
+        ) = await self.filter(
+            column=columns
+        )._verify_aggregation_operation_async()
         spec = {
             "object_ids": uuids,
             "operations": [operation],
@@ -1424,10 +1502,13 @@ class SearchContext:
             raise ex
         blob = BytesIO(res.content)
         res = self._to_sumo(prototype, blob)
+        assert isinstance(res, objects.Child)
         res._blob = blob
         return res
 
-    async def aggregate_async(self, columns=None, operation=None):
+    async def aggregate_async(
+        self, columns=None, operation=None
+    ) -> objects.Child:
         length_hidden = await self.hidden.length_async()
         if length_hidden > 0:
             return await self.hidden._aggregate_async(
@@ -1438,33 +1519,55 @@ class SearchContext:
                 columns=columns, operation=operation
             )
 
-    def aggregation(self, column=None, operation=None):
+    def aggregation(self, column=None, operation=None) -> objects.Child:
         assert operation is not None
         assert column is None or isinstance(column, str)
         sc = self.filter(aggregation=operation, column=column)
         numaggs = len(sc)
         assert numaggs <= 1
         if numaggs == 1:
-            return sc[0]
-        else:
-            return self.filter(realization=True).aggregate(
-                columns=[column] if column is not None else None,
-                operation=operation,
-            )
+            agg = sc.single
+            assert isinstance(agg, objects.Child)
+            ts = agg.metadata["_sumo"]["timestamp"]
+            reals = self.filter(
+                realization=True,
+                complex={"range": {"_sumo.timestamp": {"lt": ts}}},
+            ).realizationids
+            if set(reals) == set(
+                agg.metadata["fmu"]["aggregation"]["realization_ids"]
+            ):
+                return agg
+        # ELSE
+        return self.filter(realization=True).aggregate(
+            columns=[column] if column is not None else None,
+            operation=operation,
+        )
 
-    async def aggregation_async(self, column=None, operation=None):
+    async def aggregation_async(
+        self, column=None, operation=None
+    ) -> objects.Child:
         assert operation is not None
         assert column is None or isinstance(column, str)
         sc = self.filter(aggregation=operation, column=column)
         numaggs = await sc.length_async()
         assert numaggs <= 1
         if numaggs == 1:
-            return await sc.getitem_async(0)
-        else:
-            return await self.filter(realization=True).aggregate_async(
-                columns=[column] if column is not None else None,
-                operation=operation,
-            )
+            agg = await sc.single_async
+            assert isinstance(agg, objects.Child)
+            ts = agg.metadata["_sumo"]["timestamp"]
+            reals = await self.filter(
+                realization=True,
+                complex={"range": {"_sumo.timestamp": {"lt": ts}}},
+            ).realizationids_async
+            if set(reals) == set(
+                agg.metadata["fmu"]["aggregation"]["realization_ids"]
+            ):
+                return agg
+        # ELSE
+        return await self.filter(realization=True).aggregate_async(
+            columns=[column] if column is not None else None,
+            operation=operation,
+        )
 
     @deprecation.deprecated(
         details="Use the method 'aggregate' instead, with parameter 'operation'."

fmu/sumo/explorer/objects/case.py

@@ -4,11 +4,11 @@ from typing import Dict
 
 from sumo.wrapper import SumoClient
 
-from fmu.sumo.explorer.objects._document import Document
-from fmu.sumo.explorer.objects._search_context import SearchContext
+from ._document import Document
+from ._search_context import SearchContext
 
 
-def _make_overview_query(id):
+def _make_overview_query(id) -> Dict:
     return {
         "query": {"term": {"fmu.case.uuid.keyword": id}},
         "aggs": {
@@ -52,7 +52,7 @@ class Case(Document, SearchContext):
         self._iterations = None
 
     @property
-    def overview(self):
+    def overview(self) -> Dict:
         """Overview of case contents."""
 
         def extract_bucket_keys(bucket, name):