[io] feat: Add jsonl output option for JsonWriter

Author: AndreyMarkinPPC

libs/garf_io/README.md

@@ -11,7 +11,7 @@ Currently it supports writing data to the following destination:
 |------------| ---------------- | -------- |
 | `console`  | ConsoleWriter    | `page-size=10`,`format=table\|json`|
 | `csv`      | CsvWriter        | `destination-folder` |
-| `json`     | JsonWriter       | `destination-folder` |
+| `json`     | JsonWriter       | `destination-folder`,`format=json\|jsonl`|
 | `bq`       | BigQueryWriter   | `project`, `dataset`, `location`, `write-disposition` |
 | `sqldb`    | SqlAlchemyWriter | `connection-string`, `if-exists=fail\|replace\|append` |
 | `sheets`   | SheetsWriter     | `share-with`, `credentials-file`, `spreadsheet-url`, `is_append=True\|False`|

libs/garf_io/garf_io/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = '0.0.5'
+__version__ = '0.0.6'

libs/garf_io/garf_io/writers/json_writer.py

@@ -15,9 +15,9 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import os
+from typing import Literal
 
 import smart_open
 
@@ -34,15 +34,20 @@ class JsonWriter(file_writer.FileWriter):
   """
 
   def __init__(
-    self, destination_folder: str | os.PathLike = os.getcwd(), **kwargs: str
+    self,
+    destination_folder: str | os.PathLike = os.getcwd(),
+    format: Literal['json', 'jsonl'] = 'json',
+    **kwargs: str,
   ) -> None:
     """Initializes JsonWriter based on a destination_folder.
 
     Args:
       destination_folder: A local folder where JSON files are stored.
+      format: Format of json file ('json', 'jsonl').
       kwargs: Optional keyword arguments to initialize writer.
     """
     super().__init__(destination_folder=destination_folder, **kwargs)
+    self.format = format
 
   def write(self, report: garf_report.GarfReport, destination: str) -> str:
     """Writes Garf report to a JSON file.
@@ -55,11 +60,14 @@ def write(self, report: garf_report.GarfReport, destination: str) -> str:
       Base filename where data are written.
     """
     report = self.format_for_write(report)
-    destination = formatter.format_extension(destination, new_extension='.json')
+    file_extension = '.json' if self.format == 'json' else '.jsonl'
+    destination = formatter.format_extension(
+      destination, new_extension=file_extension
+    )
     self.create_dir()
     logging.debug('Writing %d rows of data to %s', len(report), destination)
     output_path = os.path.join(self.destination_folder, destination)
     with smart_open.open(output_path, 'w', encoding='utf-8') as f:
-      json.dump(report.to_list(row_type='dict'), f)
+      f.write(report.to_json(output=self.format))
     logging.debug('Writing to %s is completed', output_path)
     return f'[JSON] - at {output_path}'

libs/garf_io/tests/unit/writers/test_json_writer.py

@@ -48,6 +48,25 @@ def test_write_single_column_report_returns_correct_data(
 
     assert data == expected
 
+  def test_write_single_column_report_returns_correct_data(
+    self, single_column_data, output_folder
+  ):
+    writer = json_writer.JsonWriter(output_folder, format='jsonl')
+    output = output_folder / 'test.jsonl'
+    expected = [
+      {'column_1': 1},
+      {'column_1': 2},
+      {'column_1': 3},
+    ]
+
+    writer.write(single_column_data, output)
+
+    with open(output, 'r') as f:
+      data = list(f)
+
+    for i, json_str in enumerate(data):
+      assert json.loads(json_str) == expected[i]
+
   def test_write_multi_column_report_returns_correct_data(
     self, json_writer, sample_data, output_folder
   ):