Make add_arrays return fields to arrays

pystencils/datahandling/datahandling_interface.py

@@ -68,20 +68,27 @@ class DataHandling(ABC):
 
         >>> from pystencils.datahandling import create_data_handling
         >>> dh = create_data_handling((20, 30))
-        >>> dh.add_arrays('x, y(9)')
+        >>> x, y =dh.add_arrays('x, y(9)')
         >>> print(dh.fields)
-        {'x': x: double[20,30], 'y': y(9): double[20,30]}
-        >>> assert dh.fields['x'].shape = (20, 30)
-        >>> assert dh.fields['y'].index_shape = (9,)
+        {'x': x: double[22,32], 'y': y(9): double[22,32]}
+        >>> assert x == dh.fields['x']
+        >>> assert dh.fields['x'].shape == (22, 32)
+        >>> assert dh.fields['y'].index_shape == (9,)
 
         Args:
             description (str): String description of the fields to add
+        Returns:
+            Fields representing the just created arrays
         """
         from pystencils.field import _parse_part1
 
+        names = []
         for name, indices in _parse_part1(description):
+            names.append(name)
             self.add_array(name, values_per_cell=indices)
 
+        return (self.fields[n] for n in names)
+
     @abstractmethod
     def has_data(self, name):
         """Returns true if a field or custom data element with this name was added."""

pystencils_tests/test_datahandling.py

@@ -233,9 +233,11 @@ def test_add_arrays():
     field_description = 'x, y(9)'
 
     dh = create_data_handling(domain_size=domain_shape, default_ghost_layers=0, default_layout='numpy')
-    dh.add_arrays(field_description)
+    x_, y_ = dh.add_arrays(field_description)
 
     x, y = ps.fields(field_description + ': [3,4,5]')
 
+    assert x_ == x
+    assert y_ == y
     assert x == dh.fields['x']
     assert y == dh.fields['y']