Merge branch 'master' into iadd_isub_contains_improvement

quantumlib · Jan 21, 2018 · e754030 · e754030
2 parents 32c2288 + 4e6f09e
commit e754030
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/src/openfermion/utils/_operator_utils.py b/src/openfermion/utils/_operator_utils.py
@@ -255,22 +255,24 @@ def load_operator(file_name=None, data_directory=None):
     return operator
 
 
-def save_operator(operator, file_name=None, data_directory=None):
+def save_operator(operator, file_name=None, data_directory=None,
+                  allow_overwrite=False):
     """Save FermionOperator or QubitOperator to file.
 
     Args:
         operator: An instance of FermionOperator or QubitOperator.
         file_name: The name of the saved file.
         data_directory: Optional data directory to change from default data
                         directory specified in config file.
+        allow_overwrite: Whether to allow files to be overwritten.
 
     Raises:
         OperatorUtilsError: Not saved, file already exists.
         TypeError: Operator of invalid type.
     """
     file_path = get_file_path(file_name, data_directory)
 
-    if os.path.isfile(file_path):
+    if os.path.isfile(file_path) and not allow_overwrite:
         raise OperatorUtilsError("Not saved, file already exists.")
 
     if isinstance(operator, FermionOperator):

diff --git a/src/openfermion/utils/_operator_utils_test.py b/src/openfermion/utils/_operator_utils_test.py
@@ -147,6 +147,20 @@ def test_save_on_top_of_existing_operator_utils_error(self):
         with self.assertRaises(OperatorUtilsError):
             save_operator(self.fermion_operator, self.file_name)
 
+    def test_save_on_top_of_existing_operator_error_with_explicit_flag(self):
+        save_operator(self.fermion_operator, self.file_name)
+        with self.assertRaises(OperatorUtilsError):
+            save_operator(self.fermion_operator, self.file_name,
+                          allow_overwrite=False)
+
+    def test_overwrite_flag_save_on_top_of_existing_operator(self):
+        save_operator(self.fermion_operator, self.file_name)
+        save_operator(self.fermion_operator, self.file_name,
+                      allow_overwrite=True)
+        fermion_operator = load_operator(self.file_name)
+
+        self.assertTrue(fermion_operator.isclose(self.fermion_operator))
+
     def test_load_bad_type(self):
         with self.assertRaises(TypeError):
             load_operator('bad_type_operator')