mypy clean

[kiosk.git] / file_writer.py

diff --git a/file_writer.py b/file_writer.py
index 988d0a03ac49378c06079b6901cb29dc11fa5633..3cb2f39458fe475e99c438a5a38d0cd76572d874 100644 (file)
--- a/file_writer.py
+++ b/file_writer.py
@@ -1,44 +1,50 @@
+#!/usr/bin/env python3
+
 import constants
 import os
 
 
-def remove_tricky_unicode(x):
-    try:
-        x = x.decode("utf-8")
-        x = x.replace("\u2018", "'").replace("\u2019", "'")
-        x = x.replace("\u201c", '"').replace("\u201d", '"')
-        x = x.replace("\u2e3a", "-").replace("\u2014", "-")
-    except:
-        pass
-    return x
-
-
 class file_writer:
-    def __init__(self, filename):
-        self.full_filename = os.path.join(constants.pages_dir, filename)
-        self.f = open(self.full_filename, "wb")
-        self.xforms = [remove_tricky_unicode]
+    """Helper context to write a pages file."""
 
-    def add_xform(self, xform):
-        self.xforms.append(xform)
+    def __init__(self, filename: str, *, transformations=[]):
+        self.full_filename = os.path.join(constants.pages_dir, filename)
+        self.xforms = [file_writer.remove_tricky_unicode]
+        self.xforms.extend(transformations)
+        self.f = None
+
+    @staticmethod
+    def remove_tricky_unicode(x: str) -> str:
+        try:
+            x = x.replace("\u2018", "'").replace("\u2019", "'")
+            x = x.replace("\u201c", '"').replace("\u201d", '"')
+            x = x.replace("\u2e3a", "-").replace("\u2014", "-")
+        except:
+            pass
+        return x
 
     def write(self, data):
         for xform in self.xforms:
             data = xform(data)
         self.f.write(data.encode("utf-8"))
 
+    def __enter__(self):
+        self.f = open(self.full_filename, "wb")
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.close()
+
     def done(self):
-        self.f.close()
+        self.close()
 
     def close(self):
-        self.done()
+        self.f.close()
 
 
 # Test
 # def toupper(x):
-#    return x.upper()
+#   return x.upper()
 #
-# fw = file_writer("test")
-# fw.add_xform(toupper)
-# fw.write(u"This is a \u201ctest\u201d. \n")
-# fw.done()
+# with file_writer("test", transformations=[toupper]) as fw:
+#    fw.write(u"Another test!!")