slightly shorter code
This commit is contained in:
parent
7e56abf41c
commit
118247111b
|
@ -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')
|
||||
|
||||
|
|
|
@ -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)),
|
||||
|
|
Loading…
Reference in New Issue