Add support for setitem in inferencedata (#2374)

* add setitem support

* update changelog

* black

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### New features
 - Add optimized simultaneous ECDF confidence bands ([2368](https://github.com/arviz-devs/arviz/pull/2368))
+- Add support for setting groups with `idata[group]` ([2374](https://github.com/arviz-devs/arviz/pull/2374))
 
 ### Maintenance and fixes
 

arviz/data/inference_data.py

@@ -266,6 +266,14 @@ def __getitem__(self, key: str) -> xr.Dataset:
             raise KeyError(key)
         return getattr(self, key)
 
+    def __setitem__(self, key: str, value: xr.Dataset):
+        """Set item by key and update group list accordingly."""
+        if key.startswith(WARMUP_TAG):
+            self._groups_warmup.append(key)
+        else:
+            self._groups.append(key)
+        setattr(self, key, value)
+
     def groups(self) -> List[str]:
         """Return all groups present in InferenceData object."""
         return self._groups_all

arviz/tests/base_tests/test_data.py

@@ -895,6 +895,11 @@ def test_repr_html(self):
         assert escape(repr(idata)) in html
         xr.set_options(display_style=display_style)
 
+    def test_setitem(self, data_random):
+        data_random["new_group"] = data_random.posterior
+        assert "new_group" in data_random.groups()
+        assert hasattr(data_random, "new_group")
+
     def test_add_groups(self, data_random):
         data = np.random.normal(size=(4, 500, 8))
         idata = data_random