add serialize_method parameter in jsonline io methods

Author: supercoderhawk

pysenal/io/file.py

@@ -198,12 +198,13 @@ def read_jsonline_lazy(filename, encoding=_ENCODING_UTF8, default=None):
     file.close()
 
 
-def write_jsonline(filename, items, encoding=_ENCODING_UTF8):
+def write_jsonline(filename, items, encoding=_ENCODING_UTF8, serialize_method=None):
     """
     write items to file with json line format
     :param filename: destination file path
     :param items: items to be saved line by line
     :param encoding: file encoding
+    :param serialize_method: serialization method to process object
     :return: None
     """
     if isinstance(items, str):
@@ -216,7 +217,7 @@ def write_jsonline(filename, items, encoding=_ENCODING_UTF8):
         raise TypeError('items can\'t be iterable')
     file = open(filename, 'w', encoding=encoding)
     for item in items:
-        file.write(json.dumps(item, ensure_ascii=False) + '\n')
+        file.write(json.dumps(item, ensure_ascii=False, default=serialize_method) + '\n')
     file.close()
 
 
@@ -274,29 +275,31 @@ def append_lines(filename, lines, remove_file=False, encoding=_ENCODING_UTF8):
         append_line(filename, line, encoding)
 
 
-def append_jsonline(filename, item, encoding=_ENCODING_UTF8):
+def append_jsonline(filename, item, encoding=_ENCODING_UTF8, serialize_method=None):
     """
     append item as a line of json string to file
     :param filename: destination file
     :param item: item to be saved
     :param encoding: file encoding
+    :param serialize_method: serialization method to process object
     :return: None
     """
     with open(filename, 'a', encoding=encoding) as f:
-        f.write(json.dumps(item, ensure_ascii=False) + '\n')
+        f.write(json.dumps(item, ensure_ascii=False, default=serialize_method) + '\n')
 
 
-def append_jsonlines(filename, items, encoding=_ENCODING_UTF8):
+def append_jsonlines(filename, items, encoding=_ENCODING_UTF8, serialize_method=None):
     """
     append item as some lines of json string to file
     :param filename: destination file
     :param items: items to be saved
     :param encoding: file encoding
+    :param serialize_method: serialization method to process object
     :return: None
     """
     with open(filename, 'a', encoding=encoding) as f:
         for item in items:
-            f.write(json.dumps(item, ensure_ascii=False) + '\n')
+            f.write(json.dumps(item, ensure_ascii=False, default=serialize_method) + '\n')
 
 
 class __BaseFile(object):

tests/io/test_file.py

@@ -1,7 +1,10 @@
 # -*- coding: UTF-8 -*-
 import tempfile
 import pytest
+from json import JSONDecodeError
+from decimal import Decimal
 from pysenal.io.file import *
+from pysenal.utils import json_serialize
 from tests import TEST_DATA_DIR
 
 
@@ -86,6 +89,8 @@ def test_write_lines(example_lines):
 def test_jsonline(example_json):
     dirname = tempfile.gettempdir() + '/'
     filename = dirname + 'a.jsonl'
+    if os.path.exists(filename):
+        os.remove(filename)
     write_jsonline(filename, example_json)
     data = [json.loads(line) for line in read_lines(filename)]
     assert example_json == data
@@ -107,6 +112,14 @@ def test_jsonline(example_json):
     if os.path.exists(filename):
         os.remove(filename)
 
+    obj = {'n': Decimal('11.1'), 'm': Decimal('11')}
+    with pytest.raises(TypeError):
+        append_jsonline(filename, obj)
+    append_jsonline(filename, obj, serialize_method=json_serialize)
+    assert read_json(filename) == {'n': '11.1', 'm': '11'}
+    if os.path.exists(filename):
+        os.remove(filename)
+
 
 def test_text_file(example_lines):
     true_text = '\n'.join(example_lines)