feat: improve and update of geos-mesh utils functions

geos-mesh/src/geos/mesh/utils/arrayHelpers.py

@@ -7,9 +7,9 @@
 import numpy.typing as npt
 import pandas as pd  # type: ignore[import-untyped]
 import vtkmodules.util.numpy_support as vnp
-from typing import Optional, Union, cast
+from typing import Optional, Union, Any, cast
 from vtkmodules.util.numpy_support import vtk_to_numpy
-from vtkmodules.vtkCommonCore import vtkDataArray, vtkDoubleArray, vtkPoints
+from vtkmodules.vtkCommonCore import vtkDataArray, vtkPoints
 from vtkmodules.vtkCommonDataModel import ( vtkUnstructuredGrid, vtkFieldData, vtkMultiBlockDataSet, vtkDataSet,
                                             vtkCompositeDataSet, vtkDataObject, vtkPointData, vtkCellData,
                                             vtkDataObjectTreeIterator, vtkPolyData )
@@ -343,7 +343,7 @@ def isAttributeInObjectDataSet( object: vtkDataSet, attributeName: str, onPoints
     return bool( data.HasArray( attributeName ) )
 
 
-def getArrayInObject( object: vtkDataSet, attributeName: str, onPoints: bool ) -> npt.NDArray[ np.float64 ]:
+def getArrayInObject( object: vtkDataSet, attributeName: str, onPoints: bool ) -> npt.NDArray[ Any ]:
     """Return the numpy array corresponding to input attribute name in table.
 
     Args:
@@ -355,12 +355,51 @@ def getArrayInObject( object: vtkDataSet, attributeName: str, onPoints: bool ) -
     Returns:
         ArrayLike[float]: the array corresponding to input attribute name.
     """
-    array: vtkDoubleArray = getVtkArrayInObject( object, attributeName, onPoints )
-    nparray: npt.NDArray[ np.float64 ] = vnp.vtk_to_numpy( array )  # type: ignore[no-untyped-call]
+    array: vtkDataArray = getVtkArrayInObject( object, attributeName, onPoints )
+    nparray: npt.NDArray[ Any ] = vnp.vtk_to_numpy( array )  # type: ignore[no-untyped-call]
     return nparray
 
 
-def getVtkArrayInObject( object: vtkDataSet, attributeName: str, onPoints: bool ) -> vtkDoubleArray:
+def getVtkArrayTypeInObject( object: vtkDataSet, attributeName: str, onPoints: bool ) -> int:
+    """Return VTK type of requested array from dataset input.
+
+    Args:
+        object (PointSet or UnstructuredGrid): Input object.
+        attributeName (str): Name of the attribute.
+        onPoints (bool): True if attributes are on points, False if they are on cells.
+
+    Returns:
+        int: the type of the vtk array corresponding to input attribute name.
+    """
+    array: vtkDataArray = getVtkArrayInObject( object, attributeName, onPoints )
+    vtkArrayType: int = array.GetDataType()
+
+    return vtkArrayType
+
+
+def getVtkArrayTypeInMultiBlock( multiBlockDataSet: vtkMultiBlockDataSet, attributeName: str, onPoints: bool ) -> int:
+    """Return VTK type of requested array from multiblock dataset input, if existing.
+
+    Args:
+        multiBlockDataSet (PointSet or UnstructuredGrid): input object.
+        attributeName (str): Name of the attribute.
+        onPoints (bool): True if attributes are on points, False if they are on cells.
+
+    Returns:
+        int: Type of the requested vtk array if existing in input multiblock dataset, otherwise -1.
+    """
+    nbBlocks = multiBlockDataSet.GetNumberOfBlocks()
+    for idBlock in range( nbBlocks ):
+        object: vtkDataSet = cast( vtkDataSet, multiBlockDataSet.GetBlock( idBlock ) )
+        listAttributes: set[ str ] = getAttributeSet( object, onPoints )
+        if attributeName in listAttributes:
+            return getVtkArrayTypeInObject( object, attributeName, onPoints )
+
+    print( "The vtkMultiBlockDataSet has no attribute with the name " + attributeName + "." )
+    return -1
+
+
+def getVtkArrayInObject( object: vtkDataSet, attributeName: str, onPoints: bool ) -> vtkDataArray:
     """Return the array corresponding to input attribute name in table.
 
     Args:
@@ -370,7 +409,7 @@ def getVtkArrayInObject( object: vtkDataSet, attributeName: str, onPoints: bool
             on cells.
 
     Returns:
-        vtkDoubleArray: the vtk array corresponding to input attribute name.
+        vtkDataArray: the vtk array corresponding to input attribute name.
     """
     assert isAttributeInObject( object, attributeName, onPoints ), f"{attributeName} is not in input object."
     return object.GetPointData().GetArray( attributeName ) if onPoints else object.GetCellData().GetArray(
@@ -414,7 +453,7 @@ def getNumberOfComponentsDataSet( dataSet: vtkDataSet, attributeName: str, onPoi
     Returns:
         int: number of components.
     """
-    array: vtkDoubleArray = getVtkArrayInObject( dataSet, attributeName, onPoints )
+    array: vtkDataArray = getVtkArrayInObject( dataSet, attributeName, onPoints )
     return array.GetNumberOfComponents()
 
 
@@ -438,7 +477,7 @@ def getNumberOfComponentsMultiBlock(
     for blockIndex in elementaryBlockIndexes:
         block: vtkDataSet = cast( vtkDataSet, getBlockFromFlatIndex( dataSet, blockIndex ) )
         if isAttributeInObject( block, attributeName, onPoints ):
-            array: vtkDoubleArray = getVtkArrayInObject( block, attributeName, onPoints )
+            array: vtkDataArray = getVtkArrayInObject( block, attributeName, onPoints )
             return array.GetNumberOfComponents()
     return 0
 
@@ -482,7 +521,7 @@ def getComponentNamesDataSet( dataSet: vtkDataSet, attributeName: str, onPoints:
         tuple[str,...]: names of the components.
 
     """
-    array: vtkDoubleArray = getVtkArrayInObject( dataSet, attributeName, onPoints )
+    array: vtkDataArray = getVtkArrayInObject( dataSet, attributeName, onPoints )
     componentNames: list[ str ] = []
 
     if array.GetNumberOfComponents() > 1: