Add json-repair

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
   "litellm>=1.54",
   "termcolor~=2.0",
   "ipython~=8.0",
+  "json-repair~=0.35",
 ]
 authors = [
   { name="Mandana Vaziri", email="[email protected]" },

src/pdl/pdl_interpreter.py

@@ -16,6 +16,7 @@
 from typing import Any, Generator, Optional, Sequence, TypeVar  # noqa: E402
 
 import httpx  # noqa: E402
+import json_repair  # noqa: E402
 import litellm  # noqa: E402
 import yaml  # noqa: E402
 from jinja2 import (  # noqa: E402
@@ -1530,12 +1531,15 @@ def step_include(
         ) from exc
 
 
-def parse_result(parser: ParserType, text: str) -> Optional[dict[str, Any] | list[Any]]:
-    result: Optional[dict[str, Any] | list[Any]]
+JSONReturnType = dict[str, Any] | list[Any] | str | float | int | bool | None
+
+
+def parse_result(parser: ParserType, text: str) -> JSONReturnType:
+    result: JSONReturnType
     match parser:
         case "json":
             try:
-                result = json.loads(text)
+                result = json_repair.loads(text)  # type: ignore[reportAssignmentType]
             except Exception as exc:
                 raise PDLRuntimeParserError(
                     f"Attempted to parse ill-formed JSON: {repr(exc)}"

tests/test_parser.py

@@ -105,6 +105,29 @@ def test_code_parser1():
 
 
 def test_json_parser():
+    # Test json-repair,
+    # see https://github.com/mangiucugna/json_repair/blob/main/tests/test_json_repair.py
+    json_parser = """
+    text: |
+        The next 64 elements are:
+        ```json
+        { "key": "value" }
+        ```
+    parser: json
+    """
+    result = exec_str(json_parser)
+    assert result == {"key": "value"}
+
+    json_parser = """
+    text: |
+        {'key': 'string', 'key2': false, \"key3\": null, \"key4\": unquoted}
+    parser: json
+    """
+    result = exec_str(json_parser)
+    assert result == {"key": "string", "key2": False, "key3": None, "key4": "unquoted"}
+
+
+def test_jsonl_parser():
     jsonl_parser = """
     text: |
         { "a": 1, "b": 2}

tests/test_runtime_errors.py

@@ -40,10 +40,10 @@ def test_jinja_syntax():
     )
 
 
-def test_parser_json():
+def test_parser_jsonl():
     prog_str = """
 text: "{ x: 1 + 1 }"
-parser: json
+parser: jsonl
 """
     with pytest.raises(PDLRuntimeError) as exc:
         exec_str(prog_str)