diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index a3beab361..453c4ed2b 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -20,6 +20,8 @@ ARTIQ-5 edge timestamps are not required. See :mod:`artiq.coredevice.edge_counter` for the core device driver and :mod:`artiq.gateware.rtio.phy.edge_counter`/ :meth:`artiq.gateware.eem.DIO.add_std` for the gateware components. +* List datasets can now be efficiently appended to from experiments using + :meth:`artiq.language.environment.HasEnvironment.append_to_dataset`. * The controller manager now ignores device database entries without the ``"command"`` key set to facilitate sharing of devices between multiple masters. diff --git a/artiq/language/environment.py b/artiq/language/environment.py index 5eb4e8d81..0e1b2b918 100644 --- a/artiq/language/environment.py +++ b/artiq/language/environment.py @@ -321,6 +321,18 @@ class HasEnvironment: as ``slice(*sub_tuple)`` (multi-dimensional slicing).""" self.__dataset_mgr.mutate(key, index, value) + @rpc(flags={"async"}) + def append_to_dataset(self, key, value): + """Append a value to a dataset. + + The target dataset must be a list (i.e. support ``append()``), and must + have previously been set from this experiment. + + The broadcast/persist/archive mode of the given key remains unchanged + from when the dataset was last set. Appended values are transmitted + efficiently as incremental modifications in broadcast mode.""" + self.__dataset_mgr.append_to(key, value) + def get_dataset(self, key, default=NoDefault, archive=True): """Returns the contents of a dataset. diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py index a1555037f..d2dd628f5 100644 --- a/artiq/master/worker_db.py +++ b/artiq/master/worker_db.py @@ -136,17 +136,18 @@ class DatasetManager: elif key in self.local: del self.local[key] - def mutate(self, key, index, value): - target = None - if key in self.local: - target = self.local[key] + def _get_mutation_target(self, key): + target = self.local.get(key, None) if key in self._broadcaster.raw_view: if target is not None: assert target is self._broadcaster.raw_view[key][1] - target = self._broadcaster[key][1] + return self._broadcaster[key][1] if target is None: - raise KeyError("Cannot mutate non-existing dataset") + raise KeyError("Cannot mutate nonexistent dataset '{}'".format(key)) + return target + def mutate(self, key, index, value): + target = self._get_mutation_target(key) if isinstance(index, tuple): if isinstance(index[0], tuple): index = tuple(slice(*e) for e in index) @@ -154,6 +155,9 @@ class DatasetManager: index = slice(*index) setitem(target, index, value) + def append_to(self, key, value): + self._get_mutation_target(key).append(value) + def get(self, key, archive=False): if key in self.local: return self.local[key] diff --git a/artiq/test/test_datasets.py b/artiq/test/test_datasets.py index 710e34116..db35d7f34 100644 --- a/artiq/test/test_datasets.py +++ b/artiq/test/test_datasets.py @@ -32,6 +32,9 @@ class TestExperiment(EnvExperiment): def set(self, key, value, **kwargs): self.set_dataset(key, value, **kwargs) + def append(self, key, value): + self.append_to_dataset(key, value) + KEY = "foo" @@ -67,3 +70,35 @@ class ExperimentDatasetCase(unittest.TestCase): self.assertEqual(self.exp.get(KEY), 1) with self.assertRaises(KeyError): self.dataset_db.get(KEY) + + def test_append_local(self): + self.exp.set(KEY, []) + self.exp.append(KEY, 0) + self.assertEqual(self.exp.get(KEY), [0]) + self.exp.append(KEY, 1) + self.assertEqual(self.exp.get(KEY), [0, 1]) + + def test_append_broadcast(self): + self.exp.set(KEY, [], broadcast=True) + self.exp.append(KEY, 0) + self.assertEqual(self.dataset_db.data[KEY][1], [0]) + self.exp.append(KEY, 1) + self.assertEqual(self.dataset_db.data[KEY][1], [0, 1]) + + def test_append_array(self): + for broadcast in (True, False): + self.exp.set(KEY, [], broadcast=broadcast) + self.exp.append(KEY, []) + self.exp.append(KEY, []) + self.assertEqual(self.exp.get(KEY), [[], []]) + + def test_append_scalar_fails(self): + for broadcast in (True, False): + with self.assertRaises(AttributeError): + self.exp.set(KEY, 0, broadcast=broadcast) + self.exp.append(KEY, 1) + + def test_append_nonexistent_fails(self): + with self.assertRaises(KeyError): + self.exp.append(KEY, 0) +