slightly shorter code

This commit is contained in:
Philip Eisenlohr 2023-11-28 12:31:47 -05:00
parent 7e56abf41c
commit 118247111b
2 changed files with 9 additions and 11 deletions

View File

@ -601,17 +601,16 @@ class VTK:
""" """
dup = self.copy() dup = self.copy()
cell_data = dup.vtk_data.GetCellData() cell_data = dup.vtk_data.GetCellData()
for a in range(cell_data.GetNumberOfArrays()): if label in [cell_data.GetArrayName(a) for a in range(cell_data.GetNumberOfArrays())]:
if cell_data.GetArrayName(a) == label: dup.vtk_data.GetCellData().RemoveArray(label)
dup.vtk_data.GetCellData().RemoveArray(label) return dup
return dup
point_data = self.vtk_data.GetPointData() point_data = self.vtk_data.GetPointData()
for a in range(point_data.GetNumberOfArrays()): if label in [point_data.GetArrayName(a) for a in range(point_data.GetNumberOfArrays())]:
if point_data.GetArrayName(a) == label: dup.vtk_data.GetPointData().RemoveArray(label)
dup.vtk_data.GetPointData().RemoveArray(label) return dup
return dup
raise KeyError(f'array "{label}" not found') raise KeyError(f'array "{label}" not found')

View File

@ -203,9 +203,8 @@ class TestVTK:
def test_delete(self,default,mode): def test_delete(self,default,mode):
data = np.random.rand(default.N_cells if mode == 'cells' else default.N_points).astype(np.float32) data = np.random.rand(default.N_cells if mode == 'cells' else default.N_points).astype(np.float32)
v = default.set('D',data) v = default.set('D',data)
assert np.all(data == v.get('D')) assert (v.get('D') == data).all()
v = v.delete('D') assert v.delete('D') == default
assert v == default
@pytest.mark.parametrize('data_type,shape',[(float,(3,)), @pytest.mark.parametrize('data_type,shape',[(float,(3,)),
(float,(3,3)), (float,(3,3)),