We could simplify non-MPI Python code if we add another constructor to Dataset:
import openpmd_api as io
data = np.arange(size*size, dtype=np.double).reshape(3, 3)
# now
dataset = io.Dataset(data.dtype, data.shape)
# future
dataset = io.Dataset(data)