Merge pull request #123 from genlm/DRMacIver/validate-required-keys

Validate required keys in fail fast parser for JSON potential

Author: DRMacIver

genlm/control/potential/built_in/json.py

@@ -855,6 +855,11 @@ async def parse(self, input: Input):
         while True:
             await input.skip_whitespace()
             if await input.current_char() == "}":
+                if not self.required_keys.issubset(keys_seen):
+                    raise ParseError(
+                        "Missing keys: "
+                        + ", ".join(map(json.dumps, self.required_keys))
+                    )
                 await input.read(1)
                 break
             if not first:

tests/potential/test_json.py

@@ -1553,3 +1553,23 @@ def test_trie_adding_prefix_of_existing():
 
     assert trie.root.children["b"].prefix == "ar"
     assert trie.root.children["b"].accepting
+
+
+@pytest.mark.asyncio
+async def test_rejects_if_required_keys_are_missing():
+    parser = json_schema_parser(
+        {
+            "type": "object",
+            "properties": {"foo": {"type": "string"}, "bar": {"type": "number"}},
+            "required": ["foo"],
+        }
+    )
+
+    for incomplete in [
+        "{}",
+        '{"bar": 0.0}',
+    ]:
+        with pytest.raises(ParseError):
+            await parser.parse_string(incomplete)
+
+    assert await parser.parse_string('{"foo": "hello"}') == {"foo": "hello"}