Merge pull request #1 from taylorbell57/setattr

override Dataset __setattr__ function

.gitignore

@@ -129,3 +129,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# VS Code workspaces
+*.code-workspace

src/astraeus/dataset.py

@@ -0,0 +1,23 @@
+from typing import Any
+import xarray as xr
+
+
+class Dataset(xr.Dataset):
+	# Need to define this to avoid warning messages
+    __slots__ = ("_dataset",)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Overwrite the setattr function to allow behaviour like Dataset.name = value instead of the usual Dataset['name'] = (coords, value).
+        """
+        try:
+            object.__setattr__(self, name, value)
+        except AttributeError as e:
+            try:
+                if str(e) != "{!r} object has no attribute {!r}".format(
+                    type(self).__name__, name
+                ):
+                    raise e
+                else:
+                    self[name] = (list(self.coords), value)
+            except Exception as e2:
+                raise e2

src/astraeus/xarrayIO.py

@@ -1,5 +1,6 @@
 import numpy as np
 import xarray as xr
+from .dataset import Dataset
 
 
 def writeXR(filename, ds, verbose=True, append=False):
@@ -67,7 +68,7 @@ def readXR(filename, verbose=True):
            (filename.endswith(".h5") == False) and \
            (filename.endswith(".nc") == False):
             filename += ".h5"
-        ds = xr.open_dataset(filename, engine='h5netcdf')
+        ds = Dataset(xr.open_dataset(filename, engine='h5netcdf'))
         if verbose:
             print(f"Finished loading parameters from {filename}")
     except Exception as e:
@@ -257,7 +258,7 @@ def makeDataset(dictionary=None):
     ds: object
         Xarray Dataset
     """
-    ds = xr.Dataset(dictionary)
+    ds = Dataset(dictionary)
     return ds
 
 def concat(datasets, dim='time', data_vars='minimal', coords='minimal', compat='override'):