def test_add_attribute_works_for_cell_data(self):
# Given
sgrid = self.sgrid
s1 = numpy.ones(4)
sgrid.cell_data.scalars = s1
sgrid.cell_data.scalars.name = 's1'
src = VTKDataSource(data=sgrid)
self.assertEqual(src.cell_scalars_name, 's1')
self.assertEqual(sorted(src._cell_scalars_list), ['', 's1'])
# When
s2 = s1 + 1.0
src.add_attribute(s2, 's2', category='cell')
v1 = numpy.ones((4, 3))
src.add_attribute(v1, 'v1', category='cell')
t1 = numpy.ones((4, 9))
src.add_attribute(t1, 't1', category='cell')
# Then
self.assertEqual(src._cell_scalars_list, ['', 's1', 's2'])
self.assertTrue(
numpy.allclose(src.data.cell_data.get_array('s2').to_array(), s2)
)
self.assertEqual(src._cell_vectors_list, ['', 'v1'])
self.assertTrue(
numpy.allclose(src.data.cell_data.get_array('v1').to_array(), v1)
)
self.assertEqual(src._cell_tensors_list, ['', 't1'])
self.assertTrue(
numpy.allclose(src.data.cell_data.get_array('t1').to_array(), t1)
)
def test_remove_attribute(self):
# Given
sgrid = self.sgrid
s1 = numpy.ones(4)
sgrid.point_data.scalars = s1
sgrid.point_data.scalars.name = 's1'
src = VTKDataSource(data=sgrid)
self.assertEqual(src.point_scalars_name, 's1')
self.assertEqual(sorted(src._point_scalars_list), ['', 's1'])
# When
s2 = s1 + 1.0
src.add_attribute(s2, 's2')
# Then
self.assertEqual(sorted(src._point_scalars_list), ['', 's1', 's2'])
# When
src.remove_attribute('s2')
# Then
self.assertEqual(sorted(src._point_scalars_list), ['', 's1'])
self.assertEqual(src.data.point_data.get_array('s2'), None)
