From c28e06730bb16688e3248a1d71db73750e4b18ce Mon Sep 17 00:00:00 2001
From: Anna Schumaker <Anna@NoWheyCreamery.com>
Date: Tue, 5 Mar 2019 16:49:00 -0500
Subject: [PATCH] curds: Add helpful functions to the DataFile

There are cases, mostly during testing, where we might want to check if
a file exists or to remove it if it does.  Let's add those functions
now so we can use them.

Signed-off-by: Anna Schumaker <Anna@NoWheyCreamery.com>
---
 curds/data.py      | 11 ++++++++++-
 curds/test_data.py | 19 +++++++++++++++----
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/curds/data.py b/curds/data.py
index 005473d..74dc183 100644
--- a/curds/data.py
+++ b/curds/data.py
@@ -22,7 +22,7 @@ class DataFile:
     def __enter__(self):
         if self.mode == WRITE:
             self.file = open(self.temp, self.mode)
-        elif self.mode == READ and os.path.exists(self.path):
+        elif self.mode == READ and self.exists():
             self.file = open(self.path, self.mode)
         return self
 
@@ -34,10 +34,19 @@ class DataFile:
                 os.rename(self.temp, self.path)
         return True
 
+    def exists(self):
+        return os.path.exists(self.path)
+
     def pickle(self, obj):
         if self.file:
             pickle.dump(obj, self.file, pickle.HIGHEST_PROTOCOL)
 
+    def remove(self):
+        if not self.file and self.exists():
+            os.remove(self.path)
+            return True
+        return False
+
     def unpickle(self):
         if self.file:
             return pickle.load(self.file)
diff --git a/curds/test_data.py b/curds/test_data.py
index 63a436d..65c907d 100644
--- a/curds/test_data.py
+++ b/curds/test_data.py
@@ -25,13 +25,13 @@ class TestDataModule(unittest.TestCase):
         f = data.DataFile("test.file", data.READ)
         self.assertEqual(f.path, testing_file)
         self.assertEqual(f.temp, testing_temp)
-        self.assertFalse(os.path.exists(testing_file))
+        self.assertFalse(f.exists())
         self.assertEqual(f.mode, data.READ)
         self.assertIsNone(f.file)
 
         f = data.DataFile("test.file", data.WRITE)
         self.assertEqual(f.temp, testing_temp)
-        self.assertFalse(os.path.exists(testing_file))
+        self.assertFalse(f.exists())
         self.assertEqual(f.mode, data.WRITE)
         self.assertIsNone(f.file)
 
@@ -41,20 +41,31 @@ class TestDataModule(unittest.TestCase):
             self.assertIsNone(f.file)
             f.pickle(test)
             self.assertIsNone(f.unpickle())
+            self.assertFalse(f.remove())
+            self.assertTrue(f.exists())
 
         with data.DataFile("test.file", data.WRITE) as f:
             self.assertIsNotNone(f.file)
             self.assertEqual(f.file.name, testing_temp)
-            self.assertFalse(os.path.exists(testing_file))
+            self.assertFalse(f.exists())
             self.assertTrue(os.path.exists(testing_temp))
             f.pickle(test)
+            self.assertFalse(f.remove())
+            self.assertTrue(f.exists())
 
         self.assertIsNone(f.file)
         self.assertFalse(os.path.exists(testing_temp))
-        self.assertTrue(os.path.exists(testing_file))
+        self.assertTrue(f.exists())
 
         with data.DataFile("test.file", data.READ) as f:
             self.assertIsNotNone(f.file)
             self.assertEqual(f.file.name, testing_file)
             lst = f.unpickle()
             self.assertEqual(test, lst)
+            self.assertFalse(f.remove())
+            self.assertTrue(f.exists())
+
+        f = data.DataFile("test.file", data.READ)
+        self.assertTrue(f.remove())
+        self.assertFalse(f.exists())
+        self.assertFalse(f.remove())