v0.3 Feature Additions (#3)

* Added `getArrayList()` to `VTKGrid`.
* Added `getNX()`, `getNY()`, `getNZ()` to `VTKGrid`

tests/test_vtkgrid.py

@@ -13,8 +13,14 @@ def myVTKGrid(myVtkRectGrid) -> VTKGrid:
 
 def test_getDimensions(myVTKGrid, myGrid):
     assert myVTKGrid.getDimensions(0) == myGrid.nx()
+    assert myVTKGrid.getNX() == myGrid.nx()
+
     assert myVTKGrid.getDimensions(1) == myGrid.ny()
+    assert myVTKGrid.getNY() == myGrid.ny()
+
     assert myVTKGrid.getDimensions(2) == myGrid.nz()
+    assert myVTKGrid.getNZ() == myGrid.nz()
+
     assert myVTKGrid.getDimensions() == [myGrid.nx(),myGrid.ny(),myGrid.nz()]
 
 # For getting coordinates
@@ -67,6 +73,10 @@ def test_getExtentsZ(myVTKGrid, myGrid):
 def test_getArrayNone(myVTKGrid):
     assert myVTKGrid.getArray(None) is None
 
+# test getting list of arrays
+def test_getArrayList(myVtkRectGrid_V):
+    assert VTKGrid(myVtkRectGrid_V).getArrayList() == ["V"]
+
 # test reading scalar data
 def test_getArrayScalar(myVtkRectGrid_S,myScalarData):
     s=VTKGrid(myVtkRectGrid_S).getArray("S")

vtktonumpy/vtkgrid.py

@@ -73,6 +73,24 @@ def getDimensions(self, d=None):
         else:
             return self.dims[d]
 
+    def getNX(self) -> int:
+        """
+        $N_x$, the number of cells in $x$ direction
+        """
+        return self.dims[0]
+
+    def getNY(self) -> int:
+        """
+        $N_y$, the number of cells in $y$ direction
+        """
+        return self.dims[1]
+
+    def getNZ(self) -> int:
+        """
+        $N_z$, the number of cells in $z$ direction
+        """
+        return self.dims[2]
+
     # For getting coordinates
     def getXCoordinates(self) -> np.ndarray:
         """
@@ -173,6 +191,18 @@ def getExtentsZ(self) -> tuple:
             vtk_to_numpy(self.vtk_grid.GetZCoordinates())
         )
 
+    # For data arrays
+    def getArrayList(self):
+        """
+        list[str]: A list of cell arrays present in the VTK data
+        """
+        out = []
+        cell_data = self.vtk_grid.GetCellData()
+        for i in range(cell_data.GetNumberOfArrays()):
+            out.append(cell_data.GetArrayName(i))
+
+        return out
+
     def getArray(self, name: str) -> np.ndarray:
         r"""
         Get VTK array as a numpy array. If no such array exists, return :obj:`None`

vtktonumpy/vtkreader.py

@@ -119,13 +119,7 @@ def removeArray(self, array) -> None:
 
     def getArrayList(self):
         """
-        Returns a list of arrays present in the VTK data
-
-        Parameters:
-            none
-
-        Returns:
-            list[str]: a list of available cell arrays
+        list[str]: A list of cell arrays present in the VTK data
         """
         out = []
         for i in range(self.vtk_reader.GetNumberOfCellArrays()):