refactor: Refactor arrayHelpers functions

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

@@ -26,8 +26,8 @@
 
 The getter functions:
     - get the array of an attribute (one for dataset, one for multiblockDataset, one for the both and one for fieldData)
-    - get the component names of an attribute (one for dataset, one for multiblockDataset and one for the both)
-    - get the number of components of an attribute (one for dataset, one for multiblockDataset and one for the both)
+    - get the component names of an attribute
+    - get the number of components of an attribute
     - get the piece of an attribute (for any meshes)
     - get the values of an attribute as data frame (for polyData only)
     - get the vtk type of an attribute (one for dataset, one for multiblockDataset and one for the both)
@@ -450,70 +450,42 @@ def getAttributesFromDataSet( dataSet: vtkDataSet, piece: Piece ) -> dict[ str,
 
 
 def isAttributeInObject( mesh: Union[ vtkMultiBlockDataSet, vtkDataSet ], attributeName: str, piece: Piece ) -> bool:
-    """Check if an attribute is in the input object.
+    """Check if an attribute is in the input mesh for the given piece.
 
     Args:
         mesh (vtkMultiBlockDataSet | vtkDataSet): Input mesh.
         attributeName (str): Name of the attribute.
         piece (Piece): The piece of the attribute.
 
     Returns:
-        bool: True if the attribute is in the table, False otherwise.
+        bool: True if the attribute is on the mesh, False otherwise.
+
+    Raises:
+        TypeError: The mesh has to be inherited from vtkMultiBlockDataSet or vtkDataSet.
     """
-    if isinstance( mesh, vtkMultiBlockDataSet ):
-        return isAttributeInObjectMultiBlockDataSet( mesh, attributeName, piece )
-    elif isinstance( mesh, vtkDataSet ):
-        return isAttributeInObjectDataSet( mesh, attributeName, piece )
+    if isinstance( mesh, vtkDataSet ):
+        if piece == Piece.FIELD:
+            return bool( mesh.GetFieldData().HasArray( attributeName ) )
+        elif piece == Piece.POINTS:
+            return bool( mesh.GetPointData().HasArray( attributeName ) )
+        elif piece == Piece.CELLS:
+            return bool( mesh.GetCellData().HasArray( attributeName ) )
+        elif piece == Piece.BOTH:
+            onPoints: int = mesh.GetPointData().HasArray( attributeName )
+            onCells: int = mesh.GetCellData().HasArray( attributeName )
+            return onCells == onPoints == 1
+    elif isinstance( mesh, vtkMultiBlockDataSet ):
+        elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( mesh )
+        for blockIndex in elementaryBlockIndexes:
+            dataSet: vtkDataSet = vtkDataSet.SafeDownCast( mesh.GetDataSet( blockIndex ) )
+            if isAttributeInObject( dataSet, attributeName, piece ):
+                return True
     else:
         raise TypeError( "Input object must be a vtkDataSet or vtkMultiBlockDataSet." )
 
-
-def isAttributeInObjectMultiBlockDataSet( multiBlockDataSet: vtkMultiBlockDataSet, attributeName: str,
-                                          piece: Piece ) -> bool:
-    """Check if an attribute is in the input object.
-
-    Args:
-        multiBlockDataSet (vtkMultiBlockDataSet): Input multiBlockDataSet.
-        attributeName (str): Name of the attribute.
-        piece (Piece): The piece of the attribute.
-
-    Returns:
-        bool: True if the attribute is in the table, False otherwise.
-    """
-    elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( multiBlockDataSet )
-    for blockIndex in elementaryBlockIndexes:
-        dataSet: vtkDataSet = vtkDataSet.SafeDownCast( multiBlockDataSet.GetDataSet( blockIndex ) )
-        if isAttributeInObjectDataSet( dataSet, attributeName, piece ):
-            return True
-
     return False
 
 
-def isAttributeInObjectDataSet( dataSet: vtkDataSet, attributeName: str, piece: Piece ) -> bool:
-    """Check if an attribute is in the input object for the input piece.
-
-    Args:
-        dataSet (vtkDataSet): Input dataSet.
-        attributeName (str): Name of the attribute.
-        piece (Piece): The piece of the attribute.
-
-    Returns:
-        bool: True if the attribute is in the table, False otherwise.
-    """
-    if piece == Piece.FIELD:
-        return bool( dataSet.GetFieldData().HasArray( attributeName ) )
-    elif piece == Piece.POINTS:
-        return bool( dataSet.GetPointData().HasArray( attributeName ) )
-    elif piece == Piece.CELLS:
-        return bool( dataSet.GetCellData().HasArray( attributeName ) )
-    elif piece == Piece.BOTH:
-        onPoints: int = dataSet.GetPointData().HasArray( attributeName )
-        onCells: int = dataSet.GetCellData().HasArray( attributeName )
-        return onCells == onPoints == 1
-    else:
-        return False
-
-
 def isAttributeGlobal( multiBlockDataSet: vtkMultiBlockDataSet, attributeName: str, piece: Piece ) -> bool:
     """Check if an attribute is global in the input multiBlockDataSet.
 
@@ -528,7 +500,7 @@ def isAttributeGlobal( multiBlockDataSet: vtkMultiBlockDataSet, attributeName: s
     elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( multiBlockDataSet )
     for blockIndex in elementaryBlockIndexes:
         dataSet: vtkDataSet = vtkDataSet.SafeDownCast( multiBlockDataSet.GetDataSet( blockIndex ) )
-        if not isAttributeInObjectDataSet( dataSet, attributeName, piece ):
+        if not isAttributeInObject( dataSet, attributeName, piece ):
             return False
     return True
 
@@ -615,7 +587,9 @@ def getVtkArrayInObject( dataSet: vtkDataSet, attributeName: str, piece: Piece )
     Returns:
         vtkDataArray: The vtk array corresponding to input attribute name.
     """
-    assert isAttributeInObject( dataSet, attributeName, piece ), f"{attributeName} is not in input mesh."
+    if not isAttributeInObject( dataSet, attributeName, piece ):
+        raise AttributeError( f"{ attributeName } is not in input mesh." )
+
     dataArray: vtkDataArray
     if piece == Piece.POINTS:
         dataArray = dataSet.GetPointData().GetArray( attributeName )
@@ -635,7 +609,7 @@ def getNumberOfComponents(
     attributeName: str,
     piece: Piece,
 ) -> int:
-    """Get the number of components of attribute attributeName in dataSet.
+    """Get the number of components of the attribute 'attributeName' in the mesh for a given piece.
 
     Args:
         mesh (vtkMultiBlockDataSet | vtkCompositeDataSet | vtkDataSet): Mesh where the attribute is.
@@ -644,60 +618,39 @@ def getNumberOfComponents(
 
     Returns:
         int: Number of components.
+
+    Raises:
+        AttributeError: The attribute 'attributeName' is not in the mesh for the given piece.
+        TypeError: The mesh has to be inherited from vtkMultiBlockDataSet or vtkDataSet.
+        ValueError: The attribute's piece must be 'cells', 'points' or 'field'.
     """
+    nbComponents: int = 0
     if isinstance( mesh, vtkDataSet ):
-        return getNumberOfComponentsDataSet( mesh, attributeName, piece )
+        array: vtkDataArray = getVtkArrayInObject( mesh, attributeName, piece )
+        nbComponents = array.GetNumberOfComponents()
     elif isinstance( mesh, ( vtkMultiBlockDataSet, vtkCompositeDataSet ) ):
-        return getNumberOfComponentsMultiBlock( mesh, attributeName, piece )
+        elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( mesh )
+        for blockIndex in elementaryBlockIndexes:
+            dataSet: vtkDataSet = vtkDataSet.SafeDownCast( mesh.GetDataSet( blockIndex ) )
+            try:
+                return getNumberOfComponents( dataSet, attributeName, piece )
+            except ValueError as e:
+                raise e
+            except AttributeError:
+                continue
+        raise AttributeError( f"The attribute '{ attributeName }' is not in the mesh for the given piece { piece }.")
     else:
         raise TypeError( "The mesh has to be inherited from vtkMultiBlockDataSet or vtkDataSet." )
 
-
-def getNumberOfComponentsDataSet( dataSet: vtkDataSet, attributeName: str, piece: Piece ) -> int:
-    """Get the number of components of attribute attributeName in dataSet.
-
-    Args:
-        dataSet (vtkDataSet): DataSet where the attribute is.
-        attributeName (str): Name of the attribute.
-        piece (Piece): The piece of the attribute.
-
-    Returns:
-        int: Number of components.
-    """
-    array: vtkDataArray = getVtkArrayInObject( dataSet, attributeName, piece )
-    return array.GetNumberOfComponents()
-
-
-def getNumberOfComponentsMultiBlock(
-    multiBlockDataSet: Union[ vtkMultiBlockDataSet, vtkCompositeDataSet ],
-    attributeName: str,
-    piece: Piece,
-) -> int:
-    """Get the number of components of attribute attributeName in dataSet.
-
-    Args:
-        multiBlockDataSet (vtkMultiBlockDataSet | vtkCompositeDataSet): multi block data Set where the attribute is.
-        attributeName (str): Name of the attribute.
-        piece (Piece): The piece of the attribute.
-
-    Returns:
-        int: Number of components.
-    """
-    elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( multiBlockDataSet )
-    for blockIndex in elementaryBlockIndexes:
-        dataSet: vtkDataSet = vtkDataSet.SafeDownCast( multiBlockDataSet.GetDataSet( blockIndex ) )
-        if isAttributeInObject( dataSet, attributeName, piece ):
-            array: vtkDataArray = getVtkArrayInObject( dataSet, attributeName, piece )
-            return array.GetNumberOfComponents()
-    return 0
+    return nbComponents
 
 
 def getComponentNames(
     mesh: Union[ vtkMultiBlockDataSet, vtkCompositeDataSet, vtkDataSet, vtkDataObject ],
     attributeName: str,
     piece: Piece,
 ) -> tuple[ str, ...]:
-    """Get the name of the components of attribute attributeName in dataSet.
+    """Get the name of the components of the attribute 'attributeName' in a mesh.
 
     Args:
         mesh (vtkDataSet | vtkMultiBlockDataSet | vtkCompositeDataSet | vtkDataObject): Mesh where the attribute is.
@@ -706,57 +659,34 @@ def getComponentNames(
 
     Returns:
         tuple[str,...]: Names of the components.
+
+    Raises:
+        AttributeError: The attribute 'attributeName' is not in the mesh for the given piece.
+        TypeError: The mesh has to be inherited from vtkMultiBlockDataSet or vtkDataSet.
+        ValueError: The attribute's piece must be 'cells', 'points' or 'field'.
     """
+    componentNames: list[ str ] = []
     if isinstance( mesh, vtkDataSet ):
-        return getComponentNamesDataSet( mesh, attributeName, piece )
+        array: vtkDataArray = getVtkArrayInObject( mesh, attributeName, piece )
+        if array.GetNumberOfComponents() > 1:
+            componentNames += [ array.GetComponentName( i ) for i in range( array.GetNumberOfComponents() ) ]
     elif isinstance( mesh, ( vtkMultiBlockDataSet, vtkCompositeDataSet ) ):
-        return getComponentNamesMultiBlock( mesh, attributeName, piece )
+        elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( mesh )
+        for blockIndex in elementaryBlockIndexes:
+            dataSet: vtkDataSet = vtkDataSet.SafeDownCast( mesh.GetDataSet( blockIndex ) )
+            try:
+                return getComponentNames( dataSet, attributeName, piece )
+            except ValueError as e:
+                raise e
+            except AttributeError:
+                continue
+        raise AttributeError( f"The attribute '{ attributeName }' is not in the mesh for the given piece { piece }.")
     else:
-        raise TypeError( "Mesh type is not managed." )
-
-
-def getComponentNamesDataSet( dataSet: vtkDataSet, attributeName: str, piece: Piece ) -> tuple[ str, ...]:
-    """Get the name of the components of attribute attributeName in dataSet.
-
-    Args:
-        dataSet (vtkDataSet): DataSet where the attribute is.
-        attributeName (str): Name of the attribute.
-        piece (Piece): The piece of the attribute.
-
-    Returns:
-        tuple[str,...]: Names of the components.
-    """
-    array: vtkDataArray = getVtkArrayInObject( dataSet, attributeName, piece )
-    componentNames: list[ str ] = []
+        raise TypeError( "The mesh has to be inherited from vtkMultiBlockDataSet or vtkDataSet." )
 
-    if array.GetNumberOfComponents() > 1:
-        componentNames += [ array.GetComponentName( i ) for i in range( array.GetNumberOfComponents() ) ]
     return tuple( componentNames )
 
 
-def getComponentNamesMultiBlock(
-    multiBlockDataSet: Union[ vtkMultiBlockDataSet, vtkCompositeDataSet ],
-    attributeName: str,
-    piece: Piece,
-) -> tuple[ str, ...]:
-    """Get the name of the components of attribute in MultiBlockDataSet.
-
-    Args:
-        multiBlockDataSet (vtkMultiBlockDataSet | vtkCompositeDataSet): DataSet where the attribute is.
-        attributeName (str): Name of the attribute.
-        piece (Piece): The piece of the attribute.
-
-    Returns:
-        tuple[str,...]: Names of the components.
-    """
-    elementaryBlockIndexes: list[ int ] = getBlockElementIndexesFlatten( multiBlockDataSet )
-    for blockIndex in elementaryBlockIndexes:
-        dataSet: vtkDataSet = vtkDataSet.SafeDownCast( multiBlockDataSet.GetDataSet( blockIndex ) )
-        if isAttributeInObject( dataSet, attributeName, piece ):
-            return getComponentNamesDataSet( dataSet, attributeName, piece )
-    return ()
-
-
 def getAttributeValuesAsDF( surface: vtkPolyData,
                             attributeNames: tuple[ str, ...],
                             piece: Piece = Piece.CELLS ) -> pd.DataFrame: