refactor: _write to handle both local and remote filesystems

powersimdata/data_access/data_access.py

@@ -63,33 +63,39 @@ def write(self, filepath, data, save_local=True):
 
         :param str filepath: path to save data to, with extension either 'pkl', 'csv', or 'mat'.
         :param (*pandas.DataFrame* or *dict*) data: data to save
-        :raises ValueError: if extension is unknown.
+        :param bool save_local: whether a copy should also be saved to the local filesystem, if
+            such a filesystem is configured. Defaults to True.
         """
 
-        ext = os.path.basename(filepath).split(".")[-1]
-
         self._check_file_exists(filepath, should_exist=False)
 
         print("Writing %s" % filepath)
 
-        if save_local:
-            f = self.local_fs.openbin(filepath, mode="w")
-        else:
-            f = self.fs.openbin(filepath, mode="w")
+        self._write(self.fs, filepath, data)
+
+        if save_local and self.local_fs is not None:
+            self._write(self.local_fs, filepath, data)
 
-        if ext == "pkl":
-            pickle.dump(data, f)
-        elif ext == "csv":
-            data.to_csv(f)
-        elif ext == "mat":
-            savemat(f, data, appendmat=False)
-        else:
-            raise ValueError("Unknown extension! %s" % ext)
+    def _write(self, fs, filepath, data):
+        """Write a file to given data store.
 
-        f.close()
+        :param fs fs: pyfilesystem to which to write data
+        :param str filepath: path to save data to, with extension either 'pkl', 'csv', or 'mat'.
+        :param (*pandas.DataFrame* or *dict*) data: data to save
+        :raises ValueError: if extension is unknown.
+        """
 
-        if save_local and self.local_fs != self.fs:
-            fs2.copy.copy_file(self.local_fs, filepath, self.fs, filepath)
+        ext = os.path.basename(filepath).split(".")[-1]
+
+        with fs.openbin(filepath) as f:
+            if ext == "pkl":
+                pickle.dump(data, f)
+            elif ext == "csv":
+                data.to_csv(f)
+            elif ext == "mat":
+                savemat(f, data, appendmat=False)
+            else:
+                raise ValueError("Unknown extension! %s" % ext)
 
     def copy_from(self, file_name, from_dir):
         """Copy a file from data store to userspace.
@@ -204,7 +210,6 @@ def __init__(self, root=server_setup.LOCAL_DIR):
         super().__init__(root)
         self.description = "local machine"
         self.fs = fs2.open_fs(root)
-        self.local_fs = self.fs
 
     def copy_from(self, file_name, from_dir=None):
         """Copy a file from data store to userspace.