From 05399b10cd293c63020ea224d13d3b3dba1f11d4 Mon Sep 17 00:00:00 2001
From: Camille Bellot <80476446+cbellot000@users.noreply.github.com>
Date: Fri, 28 Jun 2024 11:35:16 +0200
Subject: [PATCH] collection get_entries_indices (#1640)

* get_entries_indices

* Apply suggestions from code review
---
 src/ansys/dpf/core/collection_base.py    | 29 ++++++++++++++++++++++++
 src/ansys/dpf/gate/collection_grpcapi.py | 11 +++++++++
 tests/test_fieldscontainer.py            | 12 ++++++++++
 3 files changed, 52 insertions(+)

diff --git a/src/ansys/dpf/core/collection_base.py b/src/ansys/dpf/core/collection_base.py
index 319eab003e..9438d319a6 100644
--- a/src/ansys/dpf/core/collection_base.py
+++ b/src/ansys/dpf/core/collection_base.py
@@ -26,6 +26,8 @@
 )
 from typing import Optional, Generic, TypeVar
 
+from ansys.dpf.gate.integral_types import MutableListInt32
+
 TYPE = TypeVar('TYPE')
 
 
@@ -270,6 +272,33 @@ def _get_entries(self, label_space_or_index):
                 self._api.collection_get_obj_by_index(self, label_space_or_index)
             )
 
+    @version_requires("9.0")
+    def get_entries_indices(self, label_space):
+        """Retrieve the indices of the entries corresponding a requested label space .
+
+        Notes
+        -----
+        Available starting with DPF 2025R1.
+
+        Parameters
+        ----------
+        label_space : dict[str,int]
+            Label space or index. For example,
+            ``{"time": 1, "complex": 0}`` or the index of the field.
+
+        Returns
+        -------
+        indices : list[int], list[Field], list[MeshedRegion]
+            Indices of the entries corresponding to the request.
+        """
+        client_label_space = LabelSpace(
+            label_space=label_space, obj=self, server=self._server
+        )
+        num = self._api.collection_get_num_obj_for_label_space(self, client_label_space)
+        int_list = MutableListInt32(num)
+        self._api.collection_fill_obj_indeces_for_label_space(self, client_label_space, int_list)
+        return int_list.tolist()
+
     def _get_entry(self, label_space_or_index) -> TYPE:
         """Retrieve the entry at a requested label space or index.
 
diff --git a/src/ansys/dpf/gate/collection_grpcapi.py b/src/ansys/dpf/gate/collection_grpcapi.py
index 2b3409f0e2..9f83964b44 100644
--- a/src/ansys/dpf/gate/collection_grpcapi.py
+++ b/src/ansys/dpf/gate/collection_grpcapi.py
@@ -4,6 +4,7 @@
 
 from ansys.dpf.gate.generated import collection_abstract_api
 from ansys.dpf.gate import object_handler, data_processing_grpcapi, grpc_stream_helpers, errors
+from ansys.dpf.gate.integral_types import MutableListInt32
 
 
 # -------------------------------------------------------------------------------
@@ -129,6 +130,16 @@ def _list(collection):
     def collection_get_num_obj_for_label_space(collection, space):
         return len(CollectionGRPCAPI._collection_get_entries(collection, space))
 
+    @staticmethod
+    def collection_fill_obj_indeces_for_label_space(collection, space, indices: MutableListInt32):
+        from ansys.grpc.dpf import collection_pb2
+        request = collection_pb2.EntryRequest()
+        request.collection.CopyFrom(collection._internal_obj)
+        request.label_space.CopyFrom(space._internal_obj)
+
+        out = _get_stub(collection._server).GetEntriesIndices(request)
+        indices.set(out.indices.rep_int)
+
     @staticmethod
     def collection_get_obj_by_index_for_label_space(collection, space, index):
         return data_processing_grpcapi.DataProcessingGRPCAPI.data_processing_duplicate_object_reference(
diff --git a/tests/test_fieldscontainer.py b/tests/test_fieldscontainer.py
index 92ee7964c2..15964bc0d0 100644
--- a/tests/test_fieldscontainer.py
+++ b/tests/test_fieldscontainer.py
@@ -560,3 +560,15 @@ def test_fields_container_empty_tf_support(server_type):
     fields_container = dpf.FieldsContainer(server=server_type)
 
     assert fields_container.time_freq_support == None
+
+
+@conftest.raises_for_servers_version_under("9.0")
+def test_get_entries_indices_fields_container(server_type):
+    fc = FieldsContainer(server=server_type)
+    fc.labels = ["time", "complex"]
+    for i in range(0, 20):
+        mscop = {"time": i + 1, "complex": 0}
+        fc.add_field(mscop, Field(nentities=i + 10, server=server_type))
+    assert np.allclose(fc.get_entries_indices({"time": 1, "complex": 0}), [0])
+    assert np.allclose(fc.get_entries_indices({"time": 2}), [1])
+    assert np.allclose(fc.get_entries_indices({"complex": 0}), range(0, 20))