From 118247111b75a3ac0f88e6eacc7f7bc428e629fc Mon Sep 17 00:00:00 2001 From: Philip Eisenlohr Date: Tue, 28 Nov 2023 12:31:47 -0500 Subject: [PATCH] slightly shorter code --- python/damask/_vtk.py | 15 +++++++-------- python/tests/test_VTK.py | 5 ++--- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/python/damask/_vtk.py b/python/damask/_vtk.py index 920e27781..08d8a1819 100644 --- a/python/damask/_vtk.py +++ b/python/damask/_vtk.py @@ -601,17 +601,16 @@ class VTK: """ dup = self.copy() + cell_data = dup.vtk_data.GetCellData() - for a in range(cell_data.GetNumberOfArrays()): - if cell_data.GetArrayName(a) == label: - dup.vtk_data.GetCellData().RemoveArray(label) - return dup + if label in [cell_data.GetArrayName(a) for a in range(cell_data.GetNumberOfArrays())]: + dup.vtk_data.GetCellData().RemoveArray(label) + return dup point_data = self.vtk_data.GetPointData() - for a in range(point_data.GetNumberOfArrays()): - if point_data.GetArrayName(a) == label: - dup.vtk_data.GetPointData().RemoveArray(label) - return dup + if label in [point_data.GetArrayName(a) for a in range(point_data.GetNumberOfArrays())]: + dup.vtk_data.GetPointData().RemoveArray(label) + return dup raise KeyError(f'array "{label}" not found') diff --git a/python/tests/test_VTK.py b/python/tests/test_VTK.py index e3cb64ca5..1f8237bc6 100644 --- a/python/tests/test_VTK.py +++ b/python/tests/test_VTK.py @@ -203,9 +203,8 @@ class TestVTK: def test_delete(self,default,mode): data = np.random.rand(default.N_cells if mode == 'cells' else default.N_points).astype(np.float32) v = default.set('D',data) - assert np.all(data == v.get('D')) - v = v.delete('D') - assert v == default + assert (v.get('D') == data).all() + assert v.delete('D') == default @pytest.mark.parametrize('data_type,shape',[(float,(3,)), (float,(3,3)),