fix: design and type hint of custom collections (#2695)

Author: moe-ad

src/ansys/dpf/core/__init__.py

@@ -109,17 +109,16 @@
 from ansys.dpf.core.dpf_operator import available_operator_names
 
 
-from ansys.dpf.core.collection import CollectionFactory as _CollectionFactory
 from ansys.dpf.core.collection import Collection as _Collection
 from ansys.dpf.core.label_space import LabelSpace
 
 
 # register classes for collection types:
-CustomTypeFieldsCollection:type = _CollectionFactory(CustomTypeField)
-GenericDataContainersCollection:type = _CollectionFactory(GenericDataContainer)
-StringFieldsCollection:type = _CollectionFactory(StringField)
-OperatorsCollection: type = _CollectionFactory(Operator)
-AnyCollection:type = _Collection
+CustomTypeFieldsCollection = _Collection.collection_factory(CustomTypeField)
+GenericDataContainersCollection = _Collection.collection_factory(GenericDataContainer)
+StringFieldsCollection = _Collection.collection_factory(StringField)
+OperatorsCollection = _Collection.collection_factory(Operator)
+AnyCollection = _Collection
 
 # for matplotlib
 # solves "QApplication: invalid style override passed, ignoring it."

src/ansys/dpf/core/collection.py

@@ -25,13 +25,18 @@
 
 from __future__ import annotations
 
+from typing import Generic, Type, TypeVar
+
 from ansys.dpf.core import errors, server as server_module
 from ansys.dpf.core.any import Any
-from ansys.dpf.core.collection_base import TYPE, CollectionBase
+from ansys.dpf.core.collection_base import CollectionBase
 from ansys.dpf.core.common import create_dpf_instance
 
+TYPE = TypeVar("TYPE")
+
 
-class Collection(CollectionBase[TYPE]):
+# Explicit Generic[TYPE] helps some type checkers Collection as a generic.
+class Collection(CollectionBase[TYPE], Generic[TYPE]):
     """Represents a collection of dpf objects organised by label spaces.
 
     Parameters
@@ -120,16 +125,39 @@ def add_entry(self, label_space, entry):
         """
         return super()._add_entry(label_space, Any.new_from(entry, server=self._server))
 
+    @classmethod
+    def collection_factory(cls, subtype: TYPE) -> Type[Collection[TYPE]]:
+        """Create classes deriving from Collection at runtime for a given subtype.
 
-def CollectionFactory(subtype, BaseClass=Collection):
-    """Create classes deriving from Collection at runtime for a given subtype."""
+        This factory method dynamically creates a new class that inherits from Collection
+        and is specialized for storing entries of the specified subtype.
 
-    def __init__(self, **kwargs):
-        BaseClass.__init__(self, **kwargs)
+        Parameters
+        ----------
+        subtype : type
+            Any recognized DPF type. For example, CustomTypeField, GenericDataContainer,
+            StringField, Operator, etc. This type will be used as the entries_type for
+            the new collection class.
 
-    new_class = type(
-        str(subtype.__name__) + "sCollection",
-        (BaseClass,),
-        {"__init__": __init__, "entries_type": subtype},
-    )
-    return new_class
+        Returns
+        -------
+        Type[Collection[TYPE]]
+            A new class that inherits from Collection and is specialized for the given
+            subtype. The class name will be "{subtype.__name__}sCollection".
+
+        Examples
+        --------
+        >>> from ansys.dpf.core.string_field import StringField
+        >>> from ansys.dpf.core.collection import Collection
+        >>> string_fields_collection = Collection.collection_factory(StringField)()
+        >>> string_fields_collection.__class__.__name__
+        'StringFieldsCollection'
+        >>> string_fields_collection.entries_type.__name__
+        'StringField'
+        """
+        new_class = type(
+            str(subtype.__name__) + "sCollection",
+            (cls,),
+            {"entries_type": subtype},
+        )
+        return new_class