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()
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')

View File

@ -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)),