From 44cbf4d35bf833e314ec0420d8ace2ce6887a179 Mon Sep 17 00:00:00 2001 From: Raymond Wiker Date: Mon, 25 Nov 2024 15:27:15 +0100 Subject: [PATCH] Make .filter() on Cases, Iterations and Realizations return a new object of the same class and not the base class (SearchContext). --- src/fmu/sumo/explorer/objects/cases.py | 5 +++++ src/fmu/sumo/explorer/objects/iterations.py | 4 ++++ src/fmu/sumo/explorer/objects/realizations.py | 4 ++++ tests/test_explorer.py | 4 ++-- 4 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/fmu/sumo/explorer/objects/cases.py b/src/fmu/sumo/explorer/objects/cases.py index a7d20c9..fa50b2e 100644 --- a/src/fmu/sumo/explorer/objects/cases.py +++ b/src/fmu/sumo/explorer/objects/cases.py @@ -13,4 +13,9 @@ def _maybe_prefetch(self, index): 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/iterations.py b/src/fmu/sumo/explorer/objects/iterations.py index dd942d9..bdee930 100644 --- a/src/fmu/sumo/explorer/objects/iterations.py +++ b/src/fmu/sumo/explorer/objects/iterations.py @@ -53,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 eea2ad8..405ab64 100644 --- a/src/fmu/sumo/explorer/objects/realizations.py +++ b/src/fmu/sumo/explorer/objects/realizations.py @@ -53,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 cd921df..01053ca 100644 --- a/tests/test_explorer.py +++ b/tests/test_explorer.py @@ -328,8 +328,8 @@ def test_seismic_case_by_uuid(explorer: Explorer, seismic_case_uuid: str): assert "SEGYTraceHeader" in channel_list def test_grids_and_properties(explorer: Explorer): - cases_with_grids = explorer.grids.cases.filter(status="keep").cases - cases_with_gridprops = explorer.grid_properties.cases.filter(status="keep").cases + 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