refactor: Use an enum to deal with the piece of attributes and not bool (#195)

* Create attribute piece enum
* Update the doc
* fix case with Attribute on both cells and points
* Update tests with piece enum

Author: RomainBaville

docs/geos-utils.rst

@@ -63,6 +63,14 @@ PhysicalConstants module
    :undoc-members:
    :show-inheritance:
 
+pieceEnum module
+------------------------------
+
+.. automodule:: geos.utils.pieceEnum
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 UnitRepository module
 -------------------------------------
 

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

@@ -5,13 +5,15 @@
 # ruff: noqa: E402 # disable Module level import not at top of file
 import os
 import pytest
-from typing import Union, Any, Tuple, Dict
+from typing import Union, Any
 import numpy as np
 import numpy.typing as npt
 
 from vtkmodules.vtkCommonDataModel import vtkDataSet, vtkMultiBlockDataSet, vtkPolyData
 from vtkmodules.vtkIOXML import vtkXMLGenericDataObjectReader
 
+from geos.utils.pieceEnum import Piece
+
 
 @pytest.fixture
 def arrayExpected( request: pytest.FixtureRequest ) -> npt.NDArray[ np.float64 ]:
@@ -198,19 +200,19 @@ def getElementMap() -> Any:
     """Get the element indexes mapping dictionary using the function _get_elementMap() between two meshes.
 
     Returns:
-        elementMap (Dict[int, npt.NDArray[np.int64]]): The cell mapping dictionary.
+        elementMap (dict[int, npt.NDArray[np.int64]]): The cell mapping dictionary.
     """
 
-    def _get_elementMap( meshFromName: str, meshToName: str, points: bool ) -> Dict[ int, npt.NDArray[ np.int64 ] ]:
+    def _get_elementMap( meshFromName: str, meshToName: str, piece: Piece ) -> dict[ int, npt.NDArray[ np.int64 ] ]:
         """Get the element indexes mapping dictionary between two meshes.
 
         Args:
             meshFromName (str): The name of the meshFrom.
             meshToName (str): The name of the meshTo.
-            points (bool): True if elements to map is points, False if it is cells.
+            piece (Piece): The element to map.
 
         Returns:
-            elementMap (Dict[int, npt.NDArray[np.int64]]): The element mapping dictionary.
+            elementMap (dict[int, npt.NDArray[np.int64]]): The element mapping dictionary.
         """
         sharedCells2D3DId: npt.NDArray[ np.int64 ] = np.array(
             [ [ 0, 0 ], [ 1, 1 ], [ 2, 2 ], [ 3, 3 ], [ 4, 4 ], [ 5, 5 ], [ 6, 6 ], [ 7, 7 ], [ 8, 8 ], [ 9, 9 ],
@@ -238,8 +240,8 @@ def _get_elementMap( meshFromName: str, meshToName: str, points: bool ) -> Dict[
         )
         sharedPoints1D2DId: npt.NDArray[ np.int64 ] = np.array( [ [ 0, 26 ] ], dtype=np.int64 )
         sharedPoints1D3DId: npt.NDArray[ np.int64 ] = np.array( [ [ 0, 475 ] ], dtype=np.int64 )
-        elementMap: Dict[ int, npt.NDArray[ np.int64 ] ] = {}
-        nbElements: Tuple[ int, int, int ] = ( 4092, 212, 11 ) if points else ( 1740, 156, 10 )
+        elementMap: dict[ int, npt.NDArray[ np.int64 ] ] = {}
+        nbElements: tuple[ int, int, int ] = ( 4092, 212, 11 ) if piece == Piece.POINTS else ( 1740, 156, 10 )
         if meshFromName == "well":
             if meshToName == "emptyWell":
                 elementMap[ 0 ] = np.array( [ [ 0, element ] for element in range( nbElements[ 2 ] ) ] )

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

@@ -10,6 +10,7 @@
 from geos.mesh.utils.arrayModifiers import transferAttributeWithElementMap
 from geos.mesh.utils.arrayHelpers import ( computeElementMapping, getAttributeSet, isAttributeGlobal )
 from geos.utils.Logger import ( Logger, getLogger )
+from geos.utils.pieceEnum import Piece
 
 __doc__ = """
 AttributeMapping is a vtk filter that transfers global attributes from a source mesh to a final mesh with same
@@ -41,15 +42,15 @@
     meshTo: Union[ vtkDataSet, vtkMultiBlockDataSet ]
     attributeNames: set[ str ]
     # Optional inputs.
-    onPoints: bool  # defaults to False
+    piece: Piece  # defaults to Piece.CELLS
     speHandler: bool  # defaults to False
 
     # Instantiate the filter
     attributeMappingFilter: AttributeMapping = AttributeMapping(
         meshFrom,
         meshTo,
         attributeNames,
-        onPoints,
+        piece,
         speHandler,
     )
 
@@ -77,7 +78,7 @@ def __init__(
         meshFrom: Union[ vtkDataSet, vtkMultiBlockDataSet ],
         meshTo: Union[ vtkDataSet, vtkMultiBlockDataSet ],
         attributeNames: set[ str ],
-        onPoints: bool = False,
+        piece: Piece = Piece.CELLS,
         speHandler: bool = False,
     ) -> None:
         """Transfer global attributes from a source mesh to a final mesh.
@@ -88,22 +89,20 @@ def __init__(
             meshFrom (Union[vtkDataSet, vtkMultiBlockDataSet]): The source mesh with attributes to transfer.
             meshTo (Union[vtkDataSet, vtkMultiBlockDataSet]): The final mesh where to transfer attributes.
             attributeNames (set[str]): Names of the attributes to transfer.
-            onPoints (bool): True if attributes are on points, False if they are on cells.
-                Defaults to False.
+            piece (Piece): The piece of the attribute.
+                Defaults to Piece.CELLS.
             speHandler (bool, optional): True to use a specific handler, False to use the internal handler.
                 Defaults to False.
         """
         self.meshFrom: Union[ vtkDataSet, vtkMultiBlockDataSet ] = meshFrom
         self.meshTo: Union[ vtkDataSet, vtkMultiBlockDataSet ] = meshTo
         self.attributeNames: set[ str ] = attributeNames
-        self.onPoints: bool = onPoints
-        # TODO/refact (@RomainBaville) make it an enum
-        self.piece: str = "points" if self.onPoints else "cells"
+        self.piece: Piece = piece
 
-        # cell map
+        # Element map
         self.ElementMap: dict[ int, npt.NDArray[ np.int64 ] ] = {}
 
-        # Logger.
+        # Logger
         self.logger: Logger
         if not speHandler:
             self.logger = getLogger( loggerTitle, True )
@@ -149,43 +148,42 @@ def applyFilter( self: Self ) -> None:
         self.logger.info( f"Apply filter { self.logger.name }." )
 
         if len( self.attributeNames ) == 0:
-            raise ValueError( f"Please enter at least one { self.piece } attribute to transfer." )
+            raise ValueError( "Please enter at least one attribute to transfer." )
 
-        attributesInMeshFrom: set[ str ] = getAttributeSet( self.meshFrom, self.onPoints )
+        attributesInMeshFrom: set[ str ] = getAttributeSet( self.meshFrom, self.piece )
         wrongAttributeNames: set[ str ] = self.attributeNames.difference( attributesInMeshFrom )
         if len( wrongAttributeNames ) > 0:
-            raise AttributeError(
-                f"The { self.piece } attributes { wrongAttributeNames } are not present in the source mesh." )
+            raise AttributeError( f"The attributes { wrongAttributeNames } are not present in the source mesh." )
 
-        attributesInMeshTo: set[ str ] = getAttributeSet( self.meshTo, self.onPoints )
+        attributesInMeshTo: set[ str ] = getAttributeSet( self.meshTo, self.piece )
         attributesAlreadyInMeshTo: set[ str ] = self.attributeNames.intersection( attributesInMeshTo )
         if len( attributesAlreadyInMeshTo ) > 0:
             raise AttributeError(
-                f"The { self.piece } attributes { attributesAlreadyInMeshTo } are already present in the final mesh." )
+                f"The attributes { attributesAlreadyInMeshTo } are already present in the final mesh." )
 
         if isinstance( self.meshFrom, vtkMultiBlockDataSet ):
             partialAttributes: list[ str ] = []
             for attributeName in self.attributeNames:
-                if not isAttributeGlobal( self.meshFrom, attributeName, self.onPoints ):
+                if not isAttributeGlobal( self.meshFrom, attributeName, self.piece ):
                     partialAttributes.append( attributeName )
 
             if len( partialAttributes ) > 0:
                 raise AttributeError(
-                    f"All { self.piece } attributes to transfer must be global, { partialAttributes } are partials." )
+                    f"All attributes to transfer must be global, { partialAttributes } are partials." )
 
-        self.ElementMap = computeElementMapping( self.meshFrom, self.meshTo, self.onPoints )
+        self.ElementMap = computeElementMapping( self.meshFrom, self.meshTo, self.piece )
         sharedElement: bool = False
         for key in self.ElementMap:
             if np.any( self.ElementMap[ key ] > -1 ):
                 sharedElement = True
 
         if not sharedElement:
-            raise ValueError( f"The two meshes do not have any shared { self.piece }." )
+            raise ValueError( f"The two meshes do not have any shared { self.piece.value }." )
 
         for attributeName in self.attributeNames:
             # TODO:: Modify arrayModifiers function to raise error.
             if not transferAttributeWithElementMap( self.meshFrom, self.meshTo, self.ElementMap, attributeName,
-                                                    self.onPoints, self.logger ):
+                                                    self.piece, self.logger ):
                 raise ValueError( f"Fail to transfer the attribute { attributeName }." )
 
         # Log the output message.
@@ -197,5 +195,5 @@ def _logOutputMessage( self: Self ) -> None:
         """Create and log result messages of the filter."""
         self.logger.info( f"The filter { self.logger.name } succeeded." )
         self.logger.info(
-            f"The { self.piece } attributes { self.attributeNames } have been transferred from the source mesh to the final mesh with the { self.piece } mapping."
+            f"The attributes { self.attributeNames } have been transferred from the source mesh to the final mesh with the { self.piece.value } mapping."
         )

geos-mesh/tests/conftest.py

@@ -11,6 +11,7 @@
 import vtkmodules.util.numpy_support as vnp
 from vtkmodules.vtkCommonDataModel import vtkMultiBlockDataSet, vtkDataSet
 
+from geos.utils.pieceEnum import Piece
 from geos.utils.Logger import ( getLogger, Logger, CountWarningHandler )
 from geos.mesh.utils.arrayHelpers import ( getArrayInObject, getComponentNames, getNumberOfComponents,
                                            getVtkDataTypeInObject, isAttributeGlobal, getAttributePieceInfo,
@@ -116,9 +117,7 @@ def __init__(
         # Region attribute settings.
         self.regionName: str = regionName
         self.dictRegionValues: dict[ Any, Any ] = dictRegionValues
-        self.onPoints: Union[ None, bool ]
-        self.onBoth: bool
-        self.onPoints, self.onBoth = getAttributePieceInfo( self.mesh, self.regionName )
+        self.piece: Piece = getAttributePieceInfo( self.mesh, self.regionName )
 
         # Check if the new component have default values (information for the output message).
         self.useDefaultValue: bool = False
@@ -165,15 +164,15 @@ def applyFilter( self: Self ) -> None:
         self.logger.addHandler( self.counter )
 
         # Check the validity of the attribute region.
-        if self.onPoints is None:
+        if self.piece == Piece.NONE:
             raise AttributeError( f"The attribute { self.regionName } is not in the mesh." )
 
-        if self.onBoth:
+        if self.piece == Piece.BOTH:
             raise AttributeError(
                 f"There are two attributes named { self.regionName }, one on points and the other on cells. The region attribute must be unique."
             )
 
-        nbComponentsRegion: int = getNumberOfComponents( self.mesh, self.regionName, self.onPoints )
+        nbComponentsRegion: int = getNumberOfComponents( self.mesh, self.regionName, self.piece )
         if nbComponentsRegion != 1:
             raise AttributeError( f"The region attribute { self.regionName } has to many components, one is requires." )
 
@@ -192,11 +191,11 @@ def applyFilter( self: Self ) -> None:
         newArray: npt.NDArray[ Any ]
         if isinstance( self.mesh, vtkMultiBlockDataSet ):
             # Check if the attribute region is global.
-            if not isAttributeGlobal( self.mesh, self.regionName, self.onPoints ):
+            if not isAttributeGlobal( self.mesh, self.regionName, self.piece ):
                 raise AttributeError( f"The region attribute { self.regionName } has to be global." )
 
             validIndexes, invalidIndexes = checkValidValuesInMultiBlock( self.mesh, self.regionName, listIndexes,
-                                                                         self.onPoints )
+                                                                         self.piece )
             if len( validIndexes ) == 0:
                 if len( self.dictRegionValues ) == 0:
                     self.logger.warning( "No region index entered." )
@@ -208,7 +207,7 @@ def applyFilter( self: Self ) -> None:
                                                           self.defaultValue,
                                                           self.newAttributeName,
                                                           componentNames=self.componentNames,
-                                                          onPoints=self.onPoints,
+                                                          piece=self.piece,
                                                           logger=self.logger ):
                     raise ValueError(
                         f"Something went wrong with the creation of the attribute { self.newAttributeName }." )
@@ -223,20 +222,20 @@ def applyFilter( self: Self ) -> None:
                 for flatIdDataSet in listFlatIdDataSet:
                     dataSet: vtkDataSet = vtkDataSet.SafeDownCast( self.mesh.GetDataSet( flatIdDataSet ) )
 
-                    regionArray = getArrayInObject( dataSet, self.regionName, self.onPoints )
+                    regionArray = getArrayInObject( dataSet, self.regionName, self.piece )
                     newArray = self._createArrayFromRegionArrayWithValueMap( regionArray )
                     if not createAttribute( dataSet,
                                             newArray,
                                             self.newAttributeName,
                                             componentNames=self.componentNames,
-                                            onPoints=self.onPoints,
+                                            piece=self.piece,
                                             logger=self.logger ):
                         raise ValueError(
                             f"Something went wrong with the creation of the attribute { self.newAttributeName }." )
 
         else:
             validIndexes, invalidIndexes = checkValidValuesInDataSet( self.mesh, self.regionName, listIndexes,
-                                                                      self.onPoints )
+                                                                      self.piece )
             if len( validIndexes ) == 0:
                 if len( self.dictRegionValues ) == 0:
                     self.logger.warning( "No region index entered." )
@@ -248,7 +247,7 @@ def applyFilter( self: Self ) -> None:
                                                        self.defaultValue,
                                                        self.newAttributeName,
                                                        componentNames=self.componentNames,
-                                                       onPoints=self.onPoints,
+                                                       piece=self.piece,
                                                        logger=self.logger ):
                     raise ValueError(
                         f"Something went wrong with the creation of the attribute { self.newAttributeName }." )
@@ -258,13 +257,13 @@ def applyFilter( self: Self ) -> None:
                     self.logger.warning(
                         f"The region indexes { invalidIndexes } are not in the region attribute { self.regionName }." )
 
-                regionArray = getArrayInObject( self.mesh, self.regionName, self.onPoints )
+                regionArray = getArrayInObject( self.mesh, self.regionName, self.piece )
                 newArray = self._createArrayFromRegionArrayWithValueMap( regionArray )
                 if not createAttribute( self.mesh,
                                         newArray,
                                         self.newAttributeName,
                                         componentNames=self.componentNames,
-                                        onPoints=self.onPoints,
+                                        piece=self.piece,
                                         logger=self.logger ):
                     raise ValueError(
                         f"Something went wrong with the creation of the attribute { self.newAttributeName }." )
@@ -282,7 +281,7 @@ def _setInfoRegion( self: Self ) -> None:
         """
         # Get the numpy type from the vtk typecode.
         dictType: dict[ int, Any ] = vnp.get_vtk_to_numpy_typemap()
-        regionVtkType: int = getVtkDataTypeInObject( self.mesh, self.regionName, self.onPoints )
+        regionVtkType: int = getVtkDataTypeInObject( self.mesh, self.regionName, self.piece )
         regionNpType: type = dictType[ regionVtkType ]
 
         # Set the correct type of values and region index.
@@ -356,11 +355,10 @@ def _logOutputMessage( self: Self, trueIndexes: list[ Any ] ) -> None:
 
         # Info about the created attribute.
         # The piece where the attribute is created.
-        piece: str = "points" if self.onPoints else "cells"
-        self.logger.info( f"The new attribute { self.newAttributeName } is created on { piece }." )
+        self.logger.info( f"The new attribute { self.newAttributeName } is created on { self.piece.value }." )
 
         # The number of component and they names if multiple.
-        componentNamesCreated: tuple[ str, ...] = getComponentNames( self.mesh, self.newAttributeName, self.onPoints )
+        componentNamesCreated: tuple[ str, ...] = getComponentNames( self.mesh, self.newAttributeName, self.piece )
         if self.nbComponents > 1:
             messComponent: str = ( f"The new attribute { self.newAttributeName } has { self.nbComponents } components"
                                    f" named { componentNamesCreated }." )

geos-mesh/tests/test_arrayHelpers.py

@@ -5,6 +5,7 @@
 from typing_extensions import Self
 from typing import Union, Any
 
+from geos.utils.pieceEnum import Piece
 from geos.utils.Logger import ( Logger, getLogger )
 from geos.mesh.utils.arrayModifiers import fillPartialAttributes
 from geos.mesh.utils.arrayHelpers import getAttributePieceInfo
@@ -120,22 +121,19 @@ def applyFilter( self: Self ) -> None:
             ValueError: Error during the filling of the attribute.
         """
         self.logger.info( f"Apply filter { self.logger.name }." )
-
-        onPoints: Union[ None, bool ]
-        onBoth: bool
+        piece: Piece
         for attributeName in self.dictAttributesValues:
-            onPoints, onBoth = getAttributePieceInfo( self.multiBlockDataSet, attributeName )
-            if onPoints is None:
+            piece = getAttributePieceInfo( self.multiBlockDataSet, attributeName )
+            if piece == Piece.NONE:
                 raise AttributeError( f"The attribute { attributeName } is not in the mesh." )
-
-            if onBoth:
+            elif piece == Piece.BOTH:
                 raise AttributeError(
                     f"There is two attribute named { attributeName }, one on points and the other on cells. The attribute name must be unique."
                 )
 
             if not fillPartialAttributes( self.multiBlockDataSet,
                                           attributeName,
-                                          onPoints=onPoints,
+                                          piece=piece,
                                           listValues=self.dictAttributesValues[ attributeName ],
                                           logger=self.logger ):
                 raise ValueError( "Something went wrong with the filling of partial attributes" )