Add support for writing to user-editable variables through the API

damnit/api.py

@@ -10,7 +10,8 @@
 import plotly.io as pio
 import xarray as xr
 
-from .backend.db import BlobTypes, DamnitDB
+from .kafka import UpdateProducer
+from .backend.db import BlobTypes, ReducedData, DamnitDB
 
 
 # This is a copy of damnit.ctxsupport.ctxrunner.DataType, purely so that we can
@@ -38,6 +39,13 @@ def find_proposal(propno):
     raise FileNotFoundError("Couldn't find proposal dir for {!r}".format(propno))
 
 
+# This variable is meant to be an instance of UpdateProducer, but we lazily
+# initialize it because creating the producer because takes ~100ms. Which isn't
+# very much, but it may otherwise be created hundreds of times if used in a
+# context file so it's better to avoid it where possible.
+UPDATE_PRODUCER = None
+
+
 class VariableData:
     """Represents a variable for a single run.
 
@@ -48,7 +56,7 @@ class VariableData:
     def __init__(self, name: str, title: str,
                  proposal: int, run: int,
                  h5_path: Path, data_format_version: int,
-                 db: DamnitDB, db_only: bool):
+                 db: DamnitDB, db_only: bool, missing: bool):
         self._name = name
         self._title = title
         self._proposal = proposal
@@ -57,6 +65,7 @@ def __init__(self, name: str, title: str,
         self._data_format_version = data_format_version
         self._db = db
         self._db_only = db_only
+        self._missing = missing
 
     @property
     def name(self) -> str:
@@ -145,6 +154,39 @@ def read(self):
                 # Otherwise, return a Numpy array
                 return group["data"][()]
 
+    def write(self, value, send_update=True):
+        """Write a value to a user-editable variable.
+
+        This may throw an exception if converting `value` to the type of the
+        editable variable fails, e.g. `db[100]["number"] = "foo"` will fail if
+        the `number` variable has a numeric type.
+
+        Args:
+            send_update (bool): Whether or not to send an update after
+                writing to the database. Don't use this unless you know what
+                you're doing, it may disappear in the future.
+        """
+        if not self._db_only:
+            raise RuntimeError(f"Cannot write to variable '{self.name}', it's not a user-editable variable.")
+
+        # Convert the input
+        user_variable = self._db.get_user_variables()[self.name]
+        variable_type = user_variable.get_type_class()
+        value = variable_type.to_db_value(value)
+        if value is None:
+            raise ValueError(f"Forbidden conversion of value '{value!r}' to type '{variable_type.py_type}'")
+
+        # Write to the database
+        self._db.set_variable(self.proposal, self.run, self.name, ReducedData(value))
+
+        if send_update:
+            global UPDATE_PRODUCER
+            if UPDATE_PRODUCER is None:
+                UPDATE_PRODUCER = UpdateProducer(None)
+
+            UPDATE_PRODUCER.variable_set(self.name, self.title, variable_type.type_name,
+                                      flush=True, topic=self._db.kafka_topic)
+
     def summary(self):
         """Read the summary data for a variable.
 
@@ -204,21 +246,40 @@ def file(self) -> Path:
         """The path to the HDF5 file for the run."""
         return self._h5_path
 
-    def __getitem__(self, name):
+    def _get_variable(self, name):
         key_locs = self._key_locations()
         names_to_titles = self._var_titles()
         titles_to_names = { title: name for name, title in names_to_titles.items() }
 
-        if name not in key_locs and name not in titles_to_names:
-            raise KeyError(f"Variable data for '{name!r}' not found for p{self.proposal}, r{self.run}")
-
         if name in titles_to_names:
             name = titles_to_names[name]
 
-        return VariableData(name, names_to_titles[name],
+        missing = name not in key_locs
+        user_variables = self._db.get_user_variables()
+        if missing and name in user_variables:
+            key_locs[name] = True
+        elif missing and name not in user_variables:
+            raise KeyError(f"Variable data for '{name!r}' not found for p{self.proposal}, r{self.run}")
+
+        return VariableData(name, names_to_titles.get(name),
                             self.proposal, self.run,
                             self._h5_path, self._data_format_version,
-                            self._db, key_locs[name])
+                            self._db, key_locs[name],
+                            missing)
+
+    def __getitem__(self, name):
+        variable = self._get_variable(name)
+        if variable._missing:
+            raise KeyError(f"Variable data for '{name!r}' not found for p{self.proposal}, r{self.run}")
+
+        return variable
+
+    def __setitem__(self, name, value):
+        variable = self._get_variable(name)
+
+        # The environment variable is basically only useful for tests
+        send_update = bool(int(os.environ.get("DAMNIT_SEND_UPDATE", 1)))
+        variable.write(value, send_update)
 
     def _key_locations(self):
         # Read keys from the HDF5 file

docs/api.md

@@ -31,6 +31,15 @@ data = myvar.read()
 summary = myvar.summary()
 ```
 
+You can also write to [user-editable
+variables](gui.md#adding-user-editable-variables):
+```python
+run_vars["myvar"] = 42
+
+# An alternative style would be:
+myvar.write(42)
+```
+
 ## API reference
 
 ::: damnit.Damnit

pyproject.toml

@@ -17,6 +17,7 @@ readme = "README.md"
 dependencies = [
     "h5netcdf",
     "h5py",
+    "kafka-python-ng",
     "orjson",  # used in plotly for faster json serialization
     "pandas",
     "plotly",
@@ -27,7 +28,6 @@ dependencies = [
 backend = [
     "EXtra-data",
     "ipython",
-    "kafka-python-ng",
     "kaleido",  # used in plotly to convert figures to images
     "matplotlib",
     "numpy",

tests/test_api.py

@@ -1,6 +1,8 @@
+import os
 import subprocess
 from pathlib import Path
 from textwrap import dedent
+from unittest.mock import patch
 
 import numpy as np
 import plotly.express as px
@@ -10,6 +12,7 @@
 
 from damnit import Damnit, RunVariables
 from damnit.context import ContextFile
+from damnit.backend.user_variables import UserEditableVariable
 from .helpers import extract_mock_run
 
 
@@ -90,6 +93,10 @@ def test_variable_data(mock_db_with_data, monkeypatch):
     damnit = Damnit(db_dir)
     rv = damnit[1]
 
+    # We disable updates from VariableData.write() by default so we can use the
+    # convenient RunVariables.__setitem__() method.
+    monkeypatch.setenv("DAMNIT_SEND_UPDATE", "0")
+
     # Insert a DataSet variable
     dataset_code = """
     from damnit_ctx import Variable
@@ -130,6 +137,62 @@ def dataset(run):
     assert isinstance(fig, PlotlyFigure)
     assert fig == px.bar(x=["a", "b", "c"], y=[1, 3, 2])
 
+    # It shouldn't be possible to write to non-editable variables or the default
+    # variables from DAMNIT.
+    with pytest.raises(RuntimeError):
+        rv["dataset"] = 1
+    with pytest.raises(RuntimeError):
+        rv["start_time"] = 1
+
+    # It also shouldn't be possible to write to variables that don't exist. We
+    # have to test this because RunVariables.__setitem__() will allow creating
+    # VariableData objects for variables that are missing from the run but do
+    # exist in the database.
+    with pytest.raises(KeyError):
+        rv["blah"] = 1
+
+    # Test setting an editable value
+    db.add_user_variable(UserEditableVariable("foo", "Foo", "number"))
+    rv["foo"] = 3.14
+
+    # Test sending Kafka updates
+    foo_var = rv["foo"]
+    with patch("damnit.kafka.KafkaProducer") as kafka_prd:
+        foo_var.write(42, send_update=True)
+        kafka_prd.assert_called_once()
+
+# These are smoke tests to ensure that writing of all variable types succeed,
+# tests of special cases are above in test_variable_data.
+#
+# Note that we allow any value to be converted to strings for convenience, so
+# its `bad_input` value is None.
+@pytest.mark.parametrize("variable_name,good_input,bad_input",
+                         [("boolean", True, "foo"),
+                          ("integer", 42, "foo"),
+                          ("number", 3.14, "foo"),
+                          ("stringy", "foo", None)])
+def test_writing(variable_name, good_input, bad_input, mock_db_with_data, monkeypatch):
+    db_dir, db = mock_db_with_data
+    monkeypatch.chdir(db_dir)
+    monkeypatch.setenv("DAMNIT_SEND_UPDATE", "0")
+    damnit = Damnit(db_dir)
+    rv = damnit[1]
+
+    # There's already a `string` variable in the test context file so we can't
+    # reuse the name as the type for our editable variable.
+    variable_type = "string" if variable_name == "stringy" else variable_name
+
+    # Add the user-editable variable
+    user_var = UserEditableVariable(variable_name, variable_name.capitalize(), variable_type)
+    db.add_user_variable(user_var)
+
+    rv[variable_name] = good_input
+    assert rv[variable_name].read() == good_input
+
+    if bad_input is not None:
+        with pytest.raises(ValueError):
+            rv[variable_name] = bad_input
+
 def test_api_dependencies(venv):
     package_path = Path(__file__).parent.parent
     venv.install(package_path)