feat: add ability to record derived types

Author: Matteo-Baussart-ANSYS

src/ansys/dpf/core/common.py

@@ -32,6 +32,7 @@
 import re
 import sys
 from enum import Enum
+from typing import Dict
 
 from ansys.dpf.core.misc import module_exists
 from ansys.dpf.gate.common import locations, ProgressBarBase  # noqa: F401
@@ -433,7 +434,7 @@ def type_to_special_dpf_constructors():
 _derived_class_name_to_type = None
 
 
-def derived_class_name_to_type() -> dict[str, type]:
+def derived_class_name_to_type() -> Dict[str, type]:
     """
     Returns a mapping of derived class names to their corresponding Python classes.
 
@@ -451,6 +452,30 @@ def derived_class_name_to_type() -> dict[str, type]:
     return _derived_class_name_to_type
 
 
+def record_derived_class(class_name: str, py_class: type, overwrite: bool = False):
+    """
+    Records a new derived class in the mapping of class names to their corresponding Python classes.
+
+    This function updates the global dictionary that maps derived class names (str) to their corresponding
+    Python class objects (type). If the provided class name already exists in the dictionary, it will either
+    overwrite the existing mapping or leave it unchanged based on the `overwrite` flag.
+
+    Parameters
+    ----------
+    class_name : str
+        The name of the derived class to be recorded.
+    py_class : type
+        The Python class type corresponding to the derived class.
+    overwrite : bool, optional
+        A flag indicating whether to overwrite an existing entry for the `class_name`.
+        If `True`, the entry will be overwritten. If `False` (default), the entry will
+        not be overwritten if it already exists.
+    """
+    recorded_classes = derived_class_name_to_type()
+    if overwrite or class_name not in recorded_classes:
+        recorded_classes[class_name] = py_class
+
+
 def create_dpf_instance(type, internal_obj, server):
     spe_constructors = type_to_special_dpf_constructors()
     if type in spe_constructors:

tests/test_operator.py

@@ -33,6 +33,8 @@
 from ansys import dpf
 from ansys.dpf.core import errors
 from ansys.dpf.core import operators as ops
+from ansys.dpf.core.common import derived_class_name_to_type, record_derived_class
+from ansys.dpf.core.custom_container_base import CustomContainerBase
 from ansys.dpf.core.misc import get_ansys_path
 from ansys.dpf.core.operator_specification import Specification
 from ansys.dpf.core.workflow_topology import WorkflowTopology
@@ -1457,3 +1459,26 @@ def test_operator_get_output_derived_class(server_type):
 
     workflow_topology = workflow_to_workflow_topology_op.get_output(0, WorkflowTopology)
     assert workflow_topology
+
+
+def test_record_derived_type():
+    class TestContainer(CustomContainerBase):
+        pass
+
+    class TestContainer2(CustomContainerBase):
+        pass
+
+    class_name = "TestContainer"
+
+    derived_classes = derived_class_name_to_type()
+    assert class_name not in derived_classes
+
+    record_derived_class(class_name, TestContainer)
+    assert class_name in derived_classes
+    assert derived_classes[class_name] is TestContainer
+
+    record_derived_class(class_name, TestContainer2)
+    assert derived_classes[class_name] is TestContainer
+
+    record_derived_class(class_name, TestContainer2, overwrite=True)
+    assert derived_classes[class_name] is TestContainer2