Add filtering on fmu.context.stage... (#113)

* Support filtering on fmu.context.stage. * Don't change the value of default parameters. Co-authored-by: Raymond Wiker <[email protected]>
equinor · Jan 13, 2023 · 8ca646b · 8ca646b
1 parent 61ef40a
commit 8ca646b
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 11 deletions.
diff --git a/src/fmu/sumo/explorer/_case.py b/src/fmu/sumo/explorer/_case.py
@@ -328,6 +328,8 @@ def get_objects(
         realization_ids: List[int] = (),
         aggregations: List[str] = (),
         include_time_data: TimeData = None,
+        stages: List[str] = (),
+        terms: Dict[str, List[str]] = {},
     ):
         """
         Search for child objects in a case.
@@ -338,15 +340,19 @@ def get_objects(
             `tag_names`: list of tag names (strings)
             `time_intervals`: list of time intervals (strings)
             `iteration_ids`: list of iteration ids (integers)
-            `realization_ids`: list of realizatio ids (intergers)
+            `realization_ids`: list of realization ids (integers)
             `aggregations`: list of aggregation operations (strings)
             `include_time_data`: ALL | NO_TIMEDATA | ONLY_TIMEDATA (TimeData)
+            `stages`: list of stages (strings)
+            `terms`: map of str to list of str, for additional filtering
 
         Returns:
             `DocumentCollection` used for retrieving search results
         """
 
-        terms = {"_sumo.parent_object.keyword": [self.sumo_id]}
+        terms = terms.copy()
+
+        terms["_sumo.parent_object.keyword"] = [self.sumo_id]
         fields_exists = []
 
         if iteration_ids:
@@ -365,9 +371,10 @@ def get_objects(
             terms["time_interval"] = time_intervals
 
         if aggregations:
-            terms["fmu.aggregation.operation"] = aggregations
-        else:
-            fields_exists.append("fmu.realization.id")
+            terms["fmu.aggregation.operation.keyword"] = aggregations
+
+        if stages:
+            terms["fmu.context.stage.keyword"] = stages
 
         query = self.utils.create_elastic_query(
             object_type=object_type,

diff --git a/src/fmu/sumo/explorer/_child_object.py b/src/fmu/sumo/explorer/_child_object.py
@@ -24,11 +24,15 @@ def __init__(self, sumo_client, meta_data):
         self.tag_name = fields["tag_name"][0]
         self.sumo_id = meta_data["_id"]
         self.name = source["data"]["name"]
-        self.iteration_id = source["fmu"]["iteration"]["id"]
         self.relative_path = source["file"]["relative_path"]
         self.meta_data = source
         self.object_type = source["class"]
 
+        if "iteration" in source["fmu"]:
+            self.iteration_id = source["fmu"]["iteration"]["id"]
+        else:
+            self.iteration_id = None
+
         if "realization" in source["fmu"]:
             self.realization_id = source["fmu"]["realization"]["id"]
         else:

diff --git a/src/fmu/sumo/explorer/explorer.py b/src/fmu/sumo/explorer/explorer.py
@@ -173,6 +173,8 @@ def get_objects(
         realization_ids: List[int] = (),
         aggregations: List[str] = (),
         include_time_data: TimeData = None,
+        stages: List[str] = (),
+        terms: Dict[str, List[str]] = {},
     ):  # noqa
         """
         Search for child objects in a case.
@@ -183,14 +185,17 @@ def get_objects(
             `tag_names`: list of tag names (strings)
             `time_intervals`: list of time intervals (strings)
             `iteration_ids`: list of iteration ids (integers)
-            `realization_ids`: list of realizatio ids (intergers)
+            `realization_ids`: list of realization ids (integers)
             `aggregations`: list of aggregation operations (strings)
+            `stages`: list of stages (strings)
+            `terms`: map of str to list of str, for additional filtering
 
         Returns:
             `DocumentCollection` used for retrieving search results
         """
 
-        terms = {}
+        terms = terms.copy()
+
         fields_exists = []
 
         if iteration_ids:
@@ -212,9 +217,10 @@ def get_objects(
             terms["_sumo.parent_object.keyword"] = case_ids
 
         if aggregations:
-            terms["fmu.aggregation.operation"] = aggregations
-        else:
-            fields_exists.append("fmu.realization.id")
+            terms["fmu.aggregation.operation.keyword"] = aggregations
+
+        if stages:
+            terms["fmu.context.stage.keyword"] = stages
 
         query = self.utils.create_elastic_query(
             object_type=object_type,