From 2bb55e608b45495879c00bda5d0176f6a28380c2 Mon Sep 17 00:00:00 2001
From: Maciej Bartkowiak <maciej.bartkowiak@stfc.ac.uk>
Date: Thu, 7 Mar 2024 14:14:44 +0000
Subject: [PATCH] Change the buffer in TextFormat to BytesIO

---
 .../Formats/{ASCIIFormat.py => TextFormat.py} | 39 +++++++++++--------
 1 file changed, 22 insertions(+), 17 deletions(-)
 rename MDANSE/Src/MDANSE/Framework/Formats/{ASCIIFormat.py => TextFormat.py} (88%)

diff --git a/MDANSE/Src/MDANSE/Framework/Formats/ASCIIFormat.py b/MDANSE/Src/MDANSE/Framework/Formats/TextFormat.py
similarity index 88%
rename from MDANSE/Src/MDANSE/Framework/Formats/ASCIIFormat.py
rename to MDANSE/Src/MDANSE/Framework/Formats/TextFormat.py
index 14509995bf..d350d0f505 100644
--- a/MDANSE/Src/MDANSE/Framework/Formats/ASCIIFormat.py
+++ b/MDANSE/Src/MDANSE/Framework/Formats/TextFormat.py
@@ -16,6 +16,7 @@
 import os
 import io
 import tarfile
+import codecs
 
 import numpy as np
 
@@ -23,6 +24,11 @@
 from MDANSE.Framework.Formats.IFormat import IFormat
 
 
+def length_stringio(input: "io.BytesIO") -> int:
+    result = input.getbuffer().nbytes
+    return result
+
+
 class TextFormat(IFormat):
     """
     This class handles the writing of output variables in Text format. Each output variable is written into separate Text files which are further
@@ -53,29 +59,28 @@ def write(cls, filename, data, header=""):
 
         tf = tarfile.open(filename, "w")
 
+        if header:
+            real_buffer = io.BytesIO()
+            tempStr = codecs.getwriter("utf-8")(real_buffer)
+            for line in header:
+                tempStr.write(str(line))
+            tempStr.write("\n\n")
+            real_buffer.seek(0)
+            info = tarfile.TarInfo(name="jobinfo.txt")
+            info.size = length_stringio(real_buffer)
+            tf.addfile(tarinfo=info, fileobj=real_buffer)
+
         for var in list(data.values()):
-            tempStr = io.StringIO()
+            real_buffer = io.BytesIO()
+            tempStr = codecs.getwriter("utf-8")(real_buffer)
             tempStr.write(var.info())
             tempStr.write("\n\n")
             cls.write_data(tempStr, var, data)
-            tempStr.seek(0)
+            real_buffer.seek(0)
 
             info = tarfile.TarInfo(name="%s%s" % (var.varname, cls.extensions[0]))
-            try:
-                info.size = len(tempStr)
-            except:
-                print("Ignoring the error in TextFormat, see if it matters.")
-            tf.addfile(tarinfo=info, fileobj=tempStr)
-
-        if header:
-            tempStr = io.StringIO()
-            for line in header:
-                tempStr.write(str(line))
-            tempStr.write("\n\n")
-            tempStr.seek(0)
-            info = tarfile.TarInfo(name="jobinfo.txt")
-            info.size = tempStr.len
-            tf.addfile(tarinfo=info, fileobj=tempStr)
+            info.size = length_stringio(real_buffer)
+            tf.addfile(tarinfo=info, fileobj=real_buffer)
 
         tf.close()