Merge pull request #63 from epochpic/improve-drop-variables

Improve drop variables

Author: JoelLucaAdams

src/sdf_xarray/__init__.py

@@ -253,9 +253,18 @@ def acquire_context(self, needs_lock=True):
     def load(self):  # noqa: PLR0912, PLR0915
         # Drop any requested variables
         if self.drop_variables:
+            # Build a mapping from underscored names to real variable names
+            name_map = {_rename_with_underscore(var): var for var in self.ds.variables}
+
             for variable in self.drop_variables:
-                # TODO: nicer error handling
-                self.ds.variables.pop(variable)
+                key = _rename_with_underscore(variable)
+                original_name = name_map.get(key)
+
+                if original_name is None:
+                    raise KeyError(
+                        f"Variable '{variable}' not found (interpreted as '{key}')."
+                    )
+                self.ds.variables.pop(original_name)
 
         # These two dicts are global metadata about the run or file
         attrs = {**self.ds.header, **self.ds.run_info}
@@ -434,7 +443,7 @@ def guess_can_open(self, filename_or_obj):
 
     description = "Use .sdf files in Xarray"
 
-    url = "https://epochpic.github.io/documentation/visualising_output/python.html"
+    url = "https://epochpic.github.io/documentation/visualising_output/python_beam.html"
 
 
 class SDFPreprocess:

tests/test_basic.py

@@ -185,3 +185,44 @@ def test_3d_distribution_function():
     with xr.open_dataset(EXAMPLE_3D_DIST_FN / "0000.sdf") as df:
         distribution_function = "dist_fn_x_px_py_Electron"
         assert df[distribution_function].shape == (16, 20, 20)
+
+
+def test_drop_variables():
+    with xr.open_dataset(
+        EXAMPLE_FILES_DIR / "0000.sdf", drop_variables=["Electric_Field_Ex"]
+    ) as df:
+        assert "Electric_Field_Ex" not in df
+
+
+def test_drop_variables_multiple():
+    with xr.open_dataset(
+        EXAMPLE_FILES_DIR / "0000.sdf",
+        drop_variables=["Electric_Field_Ex", "Electric_Field_Ey"],
+    ) as df:
+        assert "Electric_Field_Ex" not in df
+        assert "Electric_Field_Ey" not in df
+
+
+def test_drop_variables_original():
+    with xr.open_dataset(
+        EXAMPLE_FILES_DIR / "0000.sdf",
+        drop_variables=["Electric_Field/Ex", "Electric_Field/Ey"],
+    ) as df:
+        assert "Electric_Field_Ex" not in df
+        assert "Electric_Field_Ey" not in df
+
+
+def test_drop_variables_mixed():
+    with xr.open_dataset(
+        EXAMPLE_FILES_DIR / "0000.sdf",
+        drop_variables=["Electric_Field/Ex", "Electric_Field_Ey"],
+    ) as df:
+        assert "Electric_Field_Ex" not in df
+        assert "Electric_Field_Ey" not in df
+
+
+def test_erroring_drop_variables():
+    with pytest.raises(KeyError):
+        xr.open_dataset(
+            EXAMPLE_FILES_DIR / "0000.sdf", drop_variables=["Electric_Field/E"]
+        )