Fixed __enter__ and row writing implementation bugs. #240 and #241

nicolay-r · Jan 2, 2022 · 880436d · 880436d
1 parent b4a0b38
commit 880436d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/arekit/contrib/networks/core/predict/tsv_writer.py b/arekit/contrib/networks/core/predict/tsv_writer.py
@@ -12,7 +12,8 @@ def __init__(self, filepath):
         self.__f = None
 
     def __write(self, params):
-        self.__f.provide("{}\n".format(self.__col_separator.join(params)))
+        line = "{}\n".format(self.__col_separator.join(params))
+        self.__f.write(line.encode())
 
     def write(self, title, contents_it):
         self.__write(title)
@@ -29,6 +30,7 @@ def write(self, title, contents_it):
     def __enter__(self):
         create_dir_if_not_exists(self.__filepath)
         self.__f = gzip.open(self.__filepath, 'wb')
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.__f.close()