allow single str for load_grid field

Author: chrishavlin

yt_xarray/accessor/accessor.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import numpy as np
 import xarray as xr
@@ -22,7 +22,7 @@ def __init__(self, xarray_obj):
 
     def load_grid(
         self,
-        fields: Optional[List[str]] = None,
+        fields: Optional[Union[str, List[str]]] = None,
         geometry: str = None,
         use_callable: bool = True,
         sel_dict: Optional[dict] = None,
@@ -35,7 +35,7 @@ def load_grid(
 
         Parameters
         ----------
-        fields : list[str]
+        fields : str, list[str]
             list of fields to include. If None, will try to use all fields
 
         geometry : str
@@ -66,6 +66,11 @@ def load_grid(
             # might as well try!
             fields = list(self._obj.data_vars)
 
+        if isinstance(fields, str):
+            fields = [
+                fields,
+            ]
+
         sel_info = _xr_to_yt.Selection(
             self._obj,
             fields=fields,

yt_xarray/tests/test_accesor.py

@@ -221,3 +221,8 @@ def test_stretched_grid(use_callable):
         cell_centers = (dims[:-1] + dims[1:]) / 2
         dimvals = np.unique(ad[("index", dim)].d)
         assert np.all(dimvals == cell_centers)
+
+def test_load_single_field(ds_xr):
+    flds = "a_new_field_0"
+    ds_yt = ds_xr.yt.load_grid(flds)
+    _ = ds_yt.all_data()[("stream", flds)]