Add object hook method to JsonfileReader

flow/record/adapter/jsonfile.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-import json
-from typing import TYPE_CHECKING, BinaryIO
+import logging
+from typing import TYPE_CHECKING, Any, BinaryIO
 
 from flow import record
 from flow.record import JsonRecordPacker
@@ -28,6 +28,9 @@
 """
 
 
+log = logging.getLogger(__name__)
+
+
 class JsonfileWriter(AbstractWriter):
     fp = None
 
@@ -65,33 +68,56 @@ def close(self) -> None:
 class JsonfileReader(AbstractReader):
     fp = None
 
-    def __init__(self, path: str | Path | BinaryIO, selector: str | None = None, **kwargs):
+    def __init__(
+        self, path: str | Path | BinaryIO, selector: str | None = None, record_name: str = "json/record", **kwargs
+    ):
         self.selector = make_selector(selector)
         self.fp = record.open_path_or_stream(path, "r")
         self.packer = JsonRecordPacker()
+        self.record_name = record_name
 
     def close(self) -> None:
         if self.fp:
             self.fp.close()
         self.fp = None
 
+    def obj_hook(self, obj: record.Record | dict) -> Any:
+        return obj
+
     def __iter__(self) -> Iterator[Record]:
         ctx = get_app_context()
         selector = self.selector
+
+        if not self.fp:
+            return
+
         for line in self.fp:
-            obj = self.packer.unpack(line)
-            if isinstance(obj, record.Record):
-                if match_record_with_context(obj, selector, ctx):
-                    yield obj
-            elif isinstance(obj, record.RecordDescriptor):
-                pass
-            else:
-                # fallback for plain jsonlines (non flow.record format)
-                jd = json.loads(line)
+            if not line or line == "\n":
+                continue
+
+            try:
+                obj = self.packer.unpack(line)
+            except Exception as e:
+                log.warning("Failed unpacking line '%s'", line)
+                log.debug("", exc_info=e)
+                continue
+
+            if isinstance(obj, record.RecordDescriptor):
+                continue
+
+            elif isinstance(obj, record.Record) and match_record_with_context(obj, selector, ctx):
+                yield self.obj_hook(obj)
+
+            elif isinstance(obj, dict):
+                obj = self.obj_hook(obj)
+
                 fields = [
-                    (fieldtype_for_value(val, "string"), key) for key, val in jd.items() if not key.startswith("_")
+                    (fieldtype_for_value(val, "string"), key) for key, val in obj.items() if not key.startswith("_")
                 ]
-                desc = record.RecordDescriptor("json/record", fields)
-                obj = desc(**jd)
+                desc = record.RecordDescriptor(self.record_name, fields)
+                obj = desc(**obj)
                 if match_record_with_context(obj, selector, ctx):
                     yield obj
+
+            else:
+                log.warning("Failed handling unpacked line '%s'", line)

tests/adapter/test_json.py → tests/adapter/test_jsonfile.py

@@ -1,19 +1,20 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import pytest
 
 from flow.record import RecordReader, RecordWriter
+from flow.record.adapter.jsonfile import JsonfileReader
 from flow.record.base import RecordDescriptor
 from tests._utils import generate_records
 
 if TYPE_CHECKING:
     from pathlib import Path
 
 
-def test_json_adapter(tmp_path: Path) -> None:
+def test_jsonfile_adapter(tmp_path: Path) -> None:
     json_file = tmp_path.joinpath("records.json")
     record_adapter_path = f"jsonfile://{json_file}"
     writer = RecordWriter(record_adapter_path)
@@ -31,7 +32,7 @@ def test_json_adapter(tmp_path: Path) -> None:
     assert nr_records == nr_received_records
 
 
-def test_json_adapter_contextmanager(tmp_path: Path) -> None:
+def test_jsonfile_adapter_contextmanager(tmp_path: Path) -> None:
     json_file = tmp_path.joinpath("records.json")
     record_adapter_path = f"jsonfile://{json_file}"
     with RecordWriter(record_adapter_path) as writer:
@@ -47,7 +48,7 @@ def test_json_adapter_contextmanager(tmp_path: Path) -> None:
         assert nr_records == nr_received_records
 
 
-def test_json_adapter_jsonlines(tmp_path: Path) -> None:
+def test_jsonfile_adapter_jsonlines(tmp_path: Path) -> None:
     json_file = tmp_path.joinpath("data.jsonl")
 
     items = [
@@ -75,7 +76,7 @@ def test_json_adapter_jsonlines(tmp_path: Path) -> None:
         "jsonfile://{json_file}?descriptors=0",
     ],
 )
-def test_json_adapter_no_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None:
+def test_jsonfile_adapter_no_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None:
     json_file = tmp_path.joinpath("records.jsonl")
     record_adapter_path = record_adapter_path.format(json_file=json_file)
 
@@ -99,7 +100,7 @@ def test_json_adapter_no_record_descriptors(tmp_path: Path, record_adapter_path:
         "jsonfile://{json_file}?descriptors=1",
     ],
 )
-def test_json_adapter_with_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None:
+def test_jsonfile_adapter_with_record_descriptors(tmp_path: Path, record_adapter_path: str) -> None:
     json_file = tmp_path.joinpath("records.jsonl")
     record_adapter_path = record_adapter_path.format(json_file=json_file)
 
@@ -120,7 +121,7 @@ def test_json_adapter_with_record_descriptors(tmp_path: Path, record_adapter_pat
     assert descriptor_seen == 2
 
 
-def test_json_command_fieldtype(tmp_path: Path) -> None:
+def test_jsonfile_command_fieldtype(tmp_path: Path) -> None:
     json_file = tmp_path.joinpath("records.json")
     record_adapter_path = f"jsonfile://{json_file}"
     writer = RecordWriter(record_adapter_path)
@@ -155,3 +156,23 @@ def test_json_command_fieldtype(tmp_path: Path) -> None:
     assert records[1].commando.args == ("bash",)
 
     assert len(records) == 3
+
+
+def test_jsonfile_object_hook(tmp_path: Path) -> None:
+    """Test if the :class:`JsonFileReader` object hook method works as intended."""
+    json_file = tmp_path.joinpath("records.json")
+    json_file.write_text('{"foo": "bar"}\n{"foo": "baz"}')
+
+    class CustomJsonfileReader(JsonfileReader):
+        def obj_hook(self, obj: dict) -> Any:
+            obj["hello"] = "world"
+            return obj
+
+    with CustomJsonfileReader(json_file) as reader:
+        record = next(iter(reader))
+        assert record.foo == "bar"
+        assert record.hello == "world"
+
+        record = next(iter(reader))
+        assert record.foo == "baz"
+        assert record.hello == "world"