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>
This commit is contained in:
Anna Schumaker 2019-03-05 16:49:00 -05:00
parent 60f56ba3b7
commit c28e06730b
2 changed files with 25 additions and 5 deletions

View File

@ -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)

View File

@ -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())