From e2bb85b854bf64790dd2f553f936a06d5e52120a Mon Sep 17 00:00:00 2001
From: Anna Schumaker <Anna@NoWheyCreamery.com>
Date: Mon, 7 Jun 2021 10:17:47 -0400
Subject: [PATCH] lib: Enable reading and writing text DataFiles

Signed-off-by: Anna Schumaker <Anna@NoWheyCreamery.com>
---
 lib/data.py      | 16 +++++++++++++---
 lib/test_data.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/lib/data.py b/lib/data.py
index 7da3bb5..36a84db 100644
--- a/lib/data.py
+++ b/lib/data.py
@@ -10,6 +10,8 @@ if os.environ.get("EMMENTAL_TESTING"):
 
 READ  = 'rb'
 WRITE = 'wb'
+READ_TEXT  = 'r'
+WRITE_TEXT = 'w'
 
 emmental_data = pathlib.Path(xdg.BaseDirectory.save_data_path(__resource))
 
@@ -21,9 +23,9 @@ class DataFile:
         self.file = None
 
     def __enter__(self):
-        if self.mode == WRITE:
+        if self.mode in [ WRITE, WRITE_TEXT ]:
             self.file = self.temp.open(self.mode)
-        elif self.mode == READ and self.path.exists():
+        elif self.mode in [ READ, READ_TEXT ] and self.path.exists():
             self.file = self.path.open(self.mode)
         return self
 
@@ -31,7 +33,7 @@ class DataFile:
         if self.file:
             self.file.flush()
             self.file.close()
-            if self.mode == WRITE:
+            if self.mode in [ WRITE, WRITE_TEXT ] and exp_type == None:
                 self.temp.replace(self.path)
             self.file = None
         return exp_type == None
@@ -43,6 +45,10 @@ class DataFile:
         if self.file:
             pickle.dump(obj, self.file, pickle.HIGHEST_PROTOCOL)
 
+    def read(self):
+        if self.file:
+            return self.file.read()
+
     def remove(self):
         if self.file == None:
             self.path.unlink(missing_ok = True)
@@ -50,3 +56,7 @@ class DataFile:
     def unpickle(self):
         if self.file:
             return pickle.load(self.file)
+
+    def write(self, text):
+        if self.file:
+            self.file.write(text)
diff --git a/lib/test_data.py b/lib/test_data.py
index 735a472..6c3967e 100644
--- a/lib/test_data.py
+++ b/lib/test_data.py
@@ -20,6 +20,8 @@ class TestDataModule(unittest.TestCase):
         self.assertTrue(testing_data.is_dir())
         self.assertEqual(data.READ, 'rb')
         self.assertEqual(data.WRITE, 'wb')
+        self.assertEqual(data.READ_TEXT, 'r')
+        self.assertEqual(data.WRITE_TEXT, 'w')
 
     def test_data_file_init(self):
         f = data.DataFile("test.file", data.READ)
@@ -68,3 +70,45 @@ class TestDataModule(unittest.TestCase):
         f = data.DataFile("test.file", data.READ)
         f.remove()
         self.assertFalse(f.exists())
+
+    def test_data_file_read_write_text(self):
+        test = "test saving a string"
+        with data.DataFile("test.file", data.READ_TEXT) as f:
+            self.assertIsNone(f.file)
+            f.write(test)
+            self.assertIsNone(f.read())
+            f.remove()
+            self.assertFalse(f.exists())
+
+        with data.DataFile("test.file", data.WRITE_TEXT) as f:
+            self.assertIsNotNone(f.file)
+            self.assertEqual(pathlib.Path(f.file.name), testing_temp)
+            self.assertFalse(f.exists())
+            self.assertTrue(testing_temp.exists())
+            f.write(test)
+            f.remove()
+            self.assertFalse(f.exists())
+
+        self.assertIsNone(f.file)
+        self.assertFalse(testing_temp.exists())
+        self.assertTrue(f.exists())
+
+        with data.DataFile("test.file", data.READ_TEXT) as f:
+            self.assertIsNotNone(f.file)
+            self.assertEqual(pathlib.Path(f.file.name), testing_file)
+            txt = f.read()
+            self.assertEqual(test, txt)
+            f.remove()
+            self.assertTrue(f.exists())
+
+        f = data.DataFile("test.file", data.READ)
+        f.remove()
+        self.assertFalse(f.exists())
+
+    def test_data_file_exception(self):
+        f = data.DataFile("test.file", data.WRITE_TEXT)
+        with self.assertRaises(Exception):
+            with f:
+                f.write("test")
+                raise Exception("Test Exception")
+        self.assertFalse(f.exists())