1
0

json_schema_to_grammar.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837
  1. #!/usr/bin/env python3
  2. from __future__ import annotations
  3. import argparse
  4. import itertools
  5. import json
  6. import re
  7. import sys
  8. from typing import Any, List, Optional, Set, Tuple, Union
  9. def _build_repetition(item_rule, min_items, max_items, separator_rule=None):
  10. if max_items == 0:
  11. return ""
  12. if min_items == 0 and max_items == 1:
  13. return f'{item_rule}?'
  14. if not separator_rule:
  15. if min_items == 1 and max_items is None:
  16. return f'{item_rule}+'
  17. elif min_items == 0 and max_items is None:
  18. return f'{item_rule}*'
  19. else:
  20. return f'{item_rule}{{{min_items},{max_items if max_items is not None else ""}}}'
  21. result = item_rule + ' ' + _build_repetition(f'({separator_rule} {item_rule})', min_items - 1 if min_items > 0 else 0, max_items - 1 if max_items is not None else None)
  22. return f'({result})?' if min_items == 0 else result
  23. def _generate_min_max_int(min_value: Optional[int], max_value: Optional[int], out: list, decimals_left: int = 16, top_level: bool = True):
  24. has_min = min_value != None
  25. has_max = max_value != None
  26. def digit_range(from_char: str, to_char: str):
  27. out.append("[")
  28. if from_char == to_char:
  29. out.append(from_char)
  30. else:
  31. out.append(from_char)
  32. out.append("-")
  33. out.append(to_char)
  34. out.append("]")
  35. def more_digits(min_digits: int, max_digits: int):
  36. out.append("[0-9]")
  37. if min_digits == max_digits and min_digits == 1:
  38. return
  39. out.append("{")
  40. out.append(str(min_digits))
  41. if max_digits != min_digits:
  42. out.append(",")
  43. if max_digits != sys.maxsize:
  44. out.append(str(max_digits))
  45. out.append("}")
  46. def uniform_range(from_str: str, to_str: str):
  47. i = 0
  48. while i < len(from_str) and from_str[i] == to_str[i]:
  49. i += 1
  50. if i > 0:
  51. out.append("\"")
  52. out.append(from_str[:i])
  53. out.append("\"")
  54. if i < len(from_str):
  55. if i > 0:
  56. out.append(" ")
  57. sub_len = len(from_str) - i - 1
  58. if sub_len > 0:
  59. from_sub = from_str[i+1:]
  60. to_sub = to_str[i+1:]
  61. sub_zeros = "0" * sub_len
  62. sub_nines = "9" * sub_len
  63. to_reached = False
  64. out.append("(")
  65. if from_sub == sub_zeros:
  66. digit_range(from_str[i], chr(ord(to_str[i]) - 1))
  67. out.append(" ")
  68. more_digits(sub_len, sub_len)
  69. else:
  70. out.append("[")
  71. out.append(from_str[i])
  72. out.append("] ")
  73. out.append("(")
  74. uniform_range(from_sub, sub_nines)
  75. out.append(")")
  76. if ord(from_str[i]) < ord(to_str[i]) - 1:
  77. out.append(" | ")
  78. if to_sub == sub_nines:
  79. digit_range(chr(ord(from_str[i]) + 1), to_str[i])
  80. to_reached = True
  81. else:
  82. digit_range(chr(ord(from_str[i]) + 1), chr(ord(to_str[i]) - 1))
  83. out.append(" ")
  84. more_digits(sub_len, sub_len)
  85. if not to_reached:
  86. out.append(" | ")
  87. digit_range(to_str[i], to_str[i])
  88. out.append(" ")
  89. uniform_range(sub_zeros, to_sub)
  90. out.append(")")
  91. else:
  92. out.append("[")
  93. out.append(from_str[i])
  94. out.append("-")
  95. out.append(to_str[i])
  96. out.append("]")
  97. if has_min and has_max:
  98. if min_value < 0 and max_value < 0:
  99. out.append("\"-\" (")
  100. _generate_min_max_int(-max_value, -min_value, out, decimals_left, top_level=True)
  101. out.append(")")
  102. return
  103. if min_value < 0:
  104. out.append("\"-\" (")
  105. _generate_min_max_int(0, -min_value, out, decimals_left, top_level=True)
  106. out.append(") | ")
  107. min_value = 0
  108. min_s = str(min_value)
  109. max_s = str(max_value)
  110. min_digits = len(min_s)
  111. max_digits = len(max_s)
  112. for digits in range(min_digits, max_digits):
  113. uniform_range(min_s, "9" * digits)
  114. min_s = "1" + "0" * digits
  115. out.append(" | ")
  116. uniform_range(min_s, max_s)
  117. return
  118. less_decimals = max(decimals_left - 1, 1)
  119. if has_min:
  120. if min_value < 0:
  121. out.append("\"-\" (")
  122. _generate_min_max_int(None, -min_value, out, decimals_left, top_level=False)
  123. out.append(") | [0] | [1-9] ")
  124. more_digits(0, decimals_left - 1)
  125. elif min_value == 0:
  126. if top_level:
  127. out.append("[0] | [1-9] ")
  128. more_digits(0, less_decimals)
  129. else:
  130. more_digits(1, decimals_left)
  131. elif min_value <= 9:
  132. c = str(min_value)
  133. range_start = '1' if top_level else '0'
  134. if c > range_start:
  135. digit_range(range_start, chr(ord(c) - 1))
  136. out.append(" ")
  137. more_digits(1, less_decimals)
  138. out.append(" | ")
  139. digit_range(c, "9")
  140. out.append(" ")
  141. more_digits(0, less_decimals)
  142. else:
  143. min_s = str(min_value)
  144. length = len(min_s)
  145. c = min_s[0]
  146. if c > "1":
  147. digit_range("1" if top_level else "0", chr(ord(c) - 1))
  148. out.append(" ")
  149. more_digits(length, less_decimals)
  150. out.append(" | ")
  151. digit_range(c, c)
  152. out.append(" (")
  153. _generate_min_max_int(int(min_s[1:]), None, out, less_decimals, top_level=False)
  154. out.append(")")
  155. if c < "9":
  156. out.append(" | ")
  157. digit_range(chr(ord(c) + 1), "9")
  158. out.append(" ")
  159. more_digits(length - 1, less_decimals)
  160. return
  161. if has_max:
  162. if max_value >= 0:
  163. if top_level:
  164. out.append("\"-\" [1-9] ")
  165. more_digits(0, less_decimals)
  166. out.append(" | ")
  167. _generate_min_max_int(0, max_value, out, decimals_left, top_level=True)
  168. else:
  169. out.append("\"-\" (")
  170. _generate_min_max_int(-max_value, None, out, decimals_left, top_level=False)
  171. out.append(")")
  172. return
  173. raise RuntimeError("At least one of min_value or max_value must be set")
  174. class BuiltinRule:
  175. def __init__(self, content: str, deps: list | None = None):
  176. self.content = content
  177. self.deps = deps or []
  178. # Constraining spaces to prevent model "running away".
  179. SPACE_RULE = '| " " | "\\n"{1,2} [ \\t]{0,20}'
  180. PRIMITIVE_RULES = {
  181. 'boolean' : BuiltinRule('("true" | "false") space', []),
  182. 'decimal-part' : BuiltinRule('[0-9]{1,16}', []),
  183. 'integral-part': BuiltinRule('[0] | [1-9] [0-9]{0,15}', []),
  184. 'number' : BuiltinRule('("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space', ['integral-part', 'decimal-part']),
  185. 'integer' : BuiltinRule('("-"? integral-part) space', ['integral-part']),
  186. 'value' : BuiltinRule('object | array | string | number | boolean | null', ['object', 'array', 'string', 'number', 'boolean', 'null']),
  187. 'object' : BuiltinRule('"{" space ( string ":" space value ("," space string ":" space value)* )? "}" space', ['string', 'value']),
  188. 'array' : BuiltinRule('"[" space ( value ("," space value)* )? "]" space', ['value']),
  189. 'uuid' : BuiltinRule(r'"\"" [0-9a-fA-F]{8} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{12} "\"" space', []),
  190. 'char' : BuiltinRule(r'[^"\\\x7F\x00-\x1F] | [\\] (["\\bfnrt] | "u" [0-9a-fA-F]{4})', []),
  191. 'string' : BuiltinRule(r'"\"" char* "\"" space', ['char']),
  192. 'null' : BuiltinRule('"null" space', []),
  193. }
  194. # TODO: support "uri", "email" string formats
  195. STRING_FORMAT_RULES = {
  196. 'date' : BuiltinRule('[0-9]{4} "-" ( "0" [1-9] | "1" [0-2] ) "-" ( \"0\" [1-9] | [1-2] [0-9] | "3" [0-1] )', []),
  197. 'time' : BuiltinRule('([01] [0-9] | "2" [0-3]) ":" [0-5] [0-9] ":" [0-5] [0-9] ( "." [0-9]{3} )? ( "Z" | ( "+" | "-" ) ( [01] [0-9] | "2" [0-3] ) ":" [0-5] [0-9] )', []),
  198. 'date-time' : BuiltinRule('date "T" time', ['date', 'time']),
  199. 'date-string' : BuiltinRule('"\\"" date "\\"" space', ['date']),
  200. 'time-string' : BuiltinRule('"\\"" time "\\"" space', ['time']),
  201. 'date-time-string': BuiltinRule('"\\"" date-time "\\"" space', ['date-time']),
  202. }
  203. DOTALL = '[\\U00000000-\\U0010FFFF]'
  204. DOT = '[^\\x0A\\x0D]'
  205. RESERVED_NAMES = set(["root", "dot", *PRIMITIVE_RULES.keys(), *STRING_FORMAT_RULES.keys()])
  206. INVALID_RULE_CHARS_RE = re.compile(r'[^a-zA-Z0-9-]+')
  207. GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"\\]')
  208. GRAMMAR_RANGE_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"\]\-\\]')
  209. GRAMMAR_LITERAL_ESCAPES = {'\r': '\\r', '\n': '\\n', '"': '\\"', '-': '\\-', ']': '\\]', '\\': '\\\\'}
  210. NON_LITERAL_SET = set('|.()[]{}*+?')
  211. ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = set('^$.[]()|{}*+?')
  212. class SchemaConverter:
  213. def __init__(self, *, prop_order, allow_fetch, dotall, raw_pattern):
  214. self._prop_order = prop_order
  215. self._allow_fetch = allow_fetch
  216. self._dotall = dotall
  217. self._raw_pattern = raw_pattern
  218. self._rules = {
  219. 'space': SPACE_RULE,
  220. }
  221. self._refs = {}
  222. self._refs_being_resolved = set()
  223. def _format_literal(self, literal):
  224. escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(
  225. lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)) or m.group(0), literal
  226. )
  227. return f'"{escaped}"'
  228. def not_literal(self, literal: str, dotall: bool = True, maybe_escaped_underscores = False) -> str:
  229. '''
  230. not_literal('a') -> '[^a]'
  231. not_literal('abc') -> '([^a] | "a" ([^b] | "b" ([^c])?)?)?'
  232. '''
  233. assert len(literal) > 0, 'Empty literal not supported'
  234. def recurse(i: int):
  235. c = literal[i]
  236. if maybe_escaped_underscores and c == '_':
  237. yield f'[^{c}\\\\]'
  238. yield ' | '
  239. yield f'"\\\\"? "{c}"'
  240. else:
  241. yield f'[^{c}]'
  242. if i < len(literal) - 1:
  243. yield ' | '
  244. yield self._format_literal(c)
  245. yield ' ('
  246. yield from recurse(i + 1)
  247. yield ')?'
  248. return ''.join(('(', *recurse(0), ')'))
  249. def _not_strings(self, strings):
  250. class TrieNode:
  251. def __init__(self):
  252. self.children = {}
  253. self.is_end_of_string = False
  254. def insert(self, string):
  255. node = self
  256. for c in string:
  257. node = node.children.setdefault(c, TrieNode())
  258. node.is_end_of_string = True
  259. trie = TrieNode()
  260. for s in strings:
  261. trie.insert(s)
  262. char_rule = self._add_primitive('char', PRIMITIVE_RULES['char'])
  263. out = ['["] ( ']
  264. def visit(node):
  265. rejects = []
  266. first = True
  267. for c in sorted(node.children.keys()):
  268. child = node.children[c]
  269. rejects.append(c)
  270. if first:
  271. first = False
  272. else:
  273. out.append(' | ')
  274. out.append(f'[{c}]')
  275. if child.children:
  276. out.append(f' (')
  277. visit(child)
  278. out.append(')')
  279. elif child.is_end_of_string:
  280. out.append(f' {char_rule}+')
  281. if node.children:
  282. if not first:
  283. out.append(' | ')
  284. out.append(f'[^"{"".join(rejects)}] {char_rule}*')
  285. visit(trie)
  286. out.append(f' ){"" if trie.is_end_of_string else "?"} ["] space')
  287. return ''.join(out)
  288. def _add_rule(self, name, rule):
  289. esc_name = INVALID_RULE_CHARS_RE.sub('-', name)
  290. if esc_name not in self._rules or self._rules[esc_name] == rule:
  291. key = esc_name
  292. else:
  293. i = 0
  294. while f'{esc_name}{i}' in self._rules and self._rules[f'{esc_name}{i}'] != rule:
  295. i += 1
  296. key = f'{esc_name}{i}'
  297. self._rules[key] = rule
  298. return key
  299. def resolve_refs(self, schema: dict, url: str):
  300. '''
  301. Resolves all $ref fields in the given schema, fetching any remote schemas,
  302. replacing $ref with absolute reference URL and populating self._refs with the
  303. respective referenced (sub)schema dictionaries.
  304. '''
  305. def visit(n: dict):
  306. if isinstance(n, list):
  307. return [visit(x) for x in n]
  308. elif isinstance(n, dict):
  309. ref = n.get('$ref')
  310. if ref is not None and ref not in self._refs:
  311. if ref.startswith('https://'):
  312. assert self._allow_fetch, 'Fetching remote schemas is not allowed (use --allow-fetch for force)'
  313. import requests
  314. frag_split = ref.split('#')
  315. base_url = frag_split[0]
  316. target = self._refs.get(base_url)
  317. if target is None:
  318. target = self.resolve_refs(requests.get(ref).json(), base_url)
  319. self._refs[base_url] = target
  320. if len(frag_split) == 1 or frag_split[-1] == '':
  321. return target
  322. elif ref.startswith('#/'):
  323. target = schema
  324. ref = f'{url}{ref}'
  325. n['$ref'] = ref
  326. else:
  327. raise ValueError(f'Unsupported ref {ref}')
  328. for sel in ref.split('#')[-1].split('/')[1:]:
  329. assert target is not None, f'Error resolving ref {ref}: {sel} not in {target}'
  330. if isinstance(target, list):
  331. try:
  332. sel_index = int(sel)
  333. except ValueError:
  334. raise ValueError(f'Error resolving ref {ref}: {sel} not in {target}')
  335. assert 0 <= sel_index < len(target), f'Error resolving ref {ref}: {sel} not in {target}'
  336. target = target[sel_index]
  337. else:
  338. assert sel in target, f'Error resolving ref {ref}: {sel} not in {target}'
  339. target = target[sel]
  340. self._refs[ref] = target
  341. else:
  342. for v in n.values():
  343. visit(v)
  344. return n
  345. return visit(schema)
  346. def _generate_union_rule(self, name, alt_schemas):
  347. return ' | '.join((
  348. self.visit(alt_schema, f'{name}{"-" if name else "alternative-"}{i}')
  349. for i, alt_schema in enumerate(alt_schemas)
  350. ))
  351. def _visit_pattern(self, pattern, name):
  352. '''
  353. Transforms a regular expression pattern into a GBNF rule.
  354. Input: https://json-schema.org/understanding-json-schema/reference/regular_expressions
  355. Output: https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md
  356. Unsupported features: negative/positive lookaheads, greedy/non-greedy modifiers.
  357. Mostly a 1:1 translation, except for {x} / {x,} / {x,y} quantifiers for which
  358. we define sub-rules to keep the output lean.
  359. '''
  360. assert pattern.startswith('^') and pattern.endswith('$'), 'Pattern must start with "^" and end with "$"'
  361. pattern = pattern[1:-1]
  362. sub_rule_ids = {}
  363. i = 0
  364. length = len(pattern)
  365. def to_rule(s: tuple[str, bool]) -> str:
  366. (txt, is_literal) = s
  367. return "\"" + txt + "\"" if is_literal else txt
  368. def transform() -> tuple[str, bool]:
  369. '''
  370. Parse a unit at index i (advancing it), and return its string representation + whether it's a literal.
  371. '''
  372. nonlocal i
  373. nonlocal pattern
  374. nonlocal sub_rule_ids
  375. start = i
  376. # For each component of this sequence, store its string representation and whether it's a literal.
  377. # We only need a flat structure here to apply repetition operators to the last item, and
  378. # to merge literals at the and (we're parsing grouped ( sequences ) recursively and don't treat '|' specially
  379. # (GBNF's syntax is luckily very close to regular expressions!)
  380. seq: list[tuple[str, bool]] = []
  381. def get_dot():
  382. if self._dotall:
  383. rule = DOTALL
  384. else:
  385. # Accept any character... except \n and \r line break chars (\x0A and \xOD)
  386. rule = DOT
  387. return self._add_rule(f'dot', rule)
  388. def join_seq():
  389. nonlocal seq
  390. ret = []
  391. for is_literal, g in itertools.groupby(seq, lambda x: x[1]):
  392. if is_literal:
  393. ret.append((''.join(x[0] for x in g), True))
  394. else:
  395. ret.extend(g)
  396. if len(ret) == 1:
  397. return ret[0]
  398. return (' '.join(to_rule(x) for x in seq), False)
  399. while i < length:
  400. c = pattern[i]
  401. if c == '.':
  402. seq.append((get_dot(), False))
  403. i += 1
  404. elif c == '(':
  405. i += 1
  406. if i < length:
  407. assert pattern[i] != '?', f'Unsupported pattern syntax "{pattern[i]}" at index {i} of /{pattern}/'
  408. seq.append((f'({to_rule(transform())})', False))
  409. elif c == ')':
  410. i += 1
  411. assert start > 0 and pattern[start-1] == '(', f'Unbalanced parentheses; start = {start}, i = {i}, pattern = {pattern}'
  412. return join_seq()
  413. elif c == '[':
  414. square_brackets = c
  415. i += 1
  416. while i < length and pattern[i] != ']':
  417. if pattern[i] == '\\':
  418. square_brackets += pattern[i:i+2]
  419. i += 2
  420. else:
  421. square_brackets += pattern[i]
  422. i += 1
  423. assert i < length, f'Unbalanced square brackets; start = {start}, i = {i}, pattern = {pattern}'
  424. square_brackets += ']'
  425. i += 1
  426. seq.append((square_brackets, False))
  427. elif c == '|':
  428. seq.append(('|', False))
  429. i += 1
  430. elif c in ('*', '+', '?'):
  431. seq[-1] = (to_rule(seq[-1]) + c, False)
  432. i += 1
  433. elif c == '{':
  434. curly_brackets = c
  435. i += 1
  436. while i < length and pattern[i] != '}':
  437. curly_brackets += pattern[i]
  438. i += 1
  439. assert i < length, f'Unbalanced curly brackets; start = {start}, i = {i}, pattern = {pattern}'
  440. curly_brackets += '}'
  441. i += 1
  442. nums = [s.strip() for s in curly_brackets[1:-1].split(',')]
  443. min_times = 0
  444. max_times = None
  445. try:
  446. if len(nums) == 1:
  447. min_times = int(nums[0])
  448. max_times = min_times
  449. else:
  450. assert len(nums) == 2
  451. min_times = int(nums[0]) if nums[0] else 0
  452. max_times = int(nums[1]) if nums[1] else None
  453. except ValueError:
  454. raise ValueError(f'Invalid quantifier {curly_brackets} in /{pattern}/')
  455. (sub, sub_is_literal) = seq[-1]
  456. if not sub_is_literal:
  457. id = sub_rule_ids.get(sub)
  458. if id is None:
  459. id = self._add_rule(f'{name}-{len(sub_rule_ids) + 1}', sub)
  460. sub_rule_ids[sub] = id
  461. sub = id
  462. seq[-1] = (_build_repetition(f'"{sub}"' if sub_is_literal else sub, min_times, max_times), False)
  463. else:
  464. literal = ''
  465. while i < length:
  466. if pattern[i] == '\\' and i < length - 1:
  467. next = pattern[i + 1]
  468. if next in ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS:
  469. i += 1
  470. literal += pattern[i]
  471. i += 1
  472. else:
  473. literal += pattern[i:i+2]
  474. i += 2
  475. elif pattern[i] == '"' and not self._raw_pattern:
  476. literal += '\\"'
  477. i += 1
  478. elif pattern[i] not in NON_LITERAL_SET and \
  479. (i == length - 1 or literal == '' or pattern[i+1] == '.' or pattern[i+1] not in NON_LITERAL_SET):
  480. literal += pattern[i]
  481. i += 1
  482. else:
  483. break
  484. if literal:
  485. seq.append((literal, True))
  486. return join_seq()
  487. return self._add_rule(
  488. name,
  489. to_rule(transform()) if self._raw_pattern \
  490. else "\"\\\"\" (" + to_rule(transform()) + ") \"\\\"\" space")
  491. def _resolve_ref(self, ref):
  492. ref_fragment = ref.split('#')[-1]
  493. ref_name = 'ref' + re.sub(r'[^a-zA-Z0-9-]+', '-', ref_fragment)
  494. if ref_name not in self._rules and ref not in self._refs_being_resolved:
  495. self._refs_being_resolved.add(ref)
  496. resolved = self._refs[ref]
  497. ref_name = self.visit(resolved, ref_name)
  498. self._refs_being_resolved.remove(ref)
  499. return ref_name
  500. def _generate_constant_rule(self, value):
  501. return self._format_literal(json.dumps(value))
  502. def visit(self, schema, name):
  503. schema_type = schema.get('type')
  504. schema_format = schema.get('format')
  505. rule_name = name + '-' if name in RESERVED_NAMES else name or 'root'
  506. if (ref := schema.get('$ref')) is not None:
  507. return self._add_rule(rule_name, self._resolve_ref(ref))
  508. elif 'oneOf' in schema or 'anyOf' in schema:
  509. return self._add_rule(rule_name, self._generate_union_rule(name, schema.get('oneOf') or schema['anyOf']))
  510. elif isinstance(schema_type, list):
  511. return self._add_rule(rule_name, self._generate_union_rule(name, [{**schema, 'type': t} for t in schema_type]))
  512. elif 'const' in schema:
  513. return self._add_rule(rule_name, self._generate_constant_rule(schema['const']) + ' space')
  514. elif 'enum' in schema:
  515. rule = '(' + ' | '.join((self._generate_constant_rule(v) for v in schema['enum'])) + ') space'
  516. return self._add_rule(rule_name, rule)
  517. elif schema_type in (None, 'object') and \
  518. ('properties' in schema or \
  519. ('additionalProperties' in schema and schema['additionalProperties'] is not True)):
  520. required = set(schema.get('required', []))
  521. properties = list(schema.get('properties', {}).items())
  522. return self._add_rule(rule_name, self._build_object_rule(properties, required, name, schema.get('additionalProperties')))
  523. elif schema_type in (None, 'object', 'string') and 'allOf' in schema:
  524. required = set()
  525. properties = []
  526. enum_sets = []
  527. hybrid_name = name
  528. def add_component(comp_schema, is_required):
  529. if (ref := comp_schema.get('$ref')) is not None:
  530. comp_schema = self._refs[ref]
  531. if 'properties' in comp_schema:
  532. for prop_name, prop_schema in comp_schema['properties'].items():
  533. properties.append((prop_name, prop_schema))
  534. if is_required:
  535. required.add(prop_name)
  536. if 'enum' in comp_schema:
  537. enum_sets.append(set(comp_schema['enum']))
  538. for t in schema['allOf']:
  539. if 'anyOf' in t:
  540. for tt in t['anyOf']:
  541. add_component(tt, is_required=False)
  542. else:
  543. add_component(t, is_required=True)
  544. if enum_sets:
  545. enum_intersection = enum_sets[0]
  546. for s in enum_sets[1:]:
  547. enum_intersection &= s
  548. if enum_intersection:
  549. rule = '(' + ' | '.join((self._generate_constant_rule(v) for v in sorted(enum_intersection))) + ') space'
  550. return self._add_rule(rule_name, rule)
  551. return self._add_rule(rule_name, self._build_object_rule(properties, required, hybrid_name, additional_properties=None))
  552. elif schema_type in (None, 'array') and ('items' in schema or 'prefixItems' in schema):
  553. items = schema.get('items') or schema['prefixItems']
  554. if isinstance(items, list):
  555. return self._add_rule(
  556. rule_name,
  557. '"[" space ' +
  558. ' "," space '.join(
  559. self.visit(item, f'{name}{"-" if name else ""}tuple-{i}')
  560. for i, item in enumerate(items)) +
  561. ' "]" space')
  562. else:
  563. item_rule_name = self.visit(items, f'{name}{"-" if name else ""}item')
  564. min_items = schema.get("minItems", 0)
  565. max_items = schema.get("maxItems")
  566. return self._add_rule(rule_name, '"[" space ' + _build_repetition(item_rule_name, min_items, max_items, separator_rule='"," space') + ' "]" space')
  567. elif schema_type in (None, 'string') and 'pattern' in schema:
  568. return self._visit_pattern(schema['pattern'], rule_name)
  569. elif schema_type in (None, 'string') and re.match(r'^uuid[1-5]?$', schema_format or ''):
  570. return self._add_primitive(
  571. 'root' if rule_name == 'root' else schema_format,
  572. PRIMITIVE_RULES['uuid']
  573. )
  574. elif schema_type in (None, 'string') and f'{schema_format}-string' in STRING_FORMAT_RULES:
  575. prim_name = f'{schema_format}-string'
  576. return self._add_rule(rule_name, self._add_primitive(prim_name, STRING_FORMAT_RULES[prim_name]))
  577. elif schema_type == 'string' and ('minLength' in schema or 'maxLength' in schema):
  578. char_rule = self._add_primitive('char', PRIMITIVE_RULES['char'])
  579. min_len = schema.get('minLength', 0)
  580. max_len = schema.get('maxLength')
  581. return self._add_rule(rule_name, r'"\"" ' + _build_repetition(char_rule, min_len, max_len) + r' "\"" space')
  582. elif schema_type in (None, 'integer') and \
  583. ('minimum' in schema or 'exclusiveMinimum' in schema or 'maximum' in schema or 'exclusiveMaximum' in schema):
  584. min_value = None
  585. max_value = None
  586. if 'minimum' in schema:
  587. min_value = schema['minimum']
  588. elif 'exclusiveMinimum' in schema:
  589. min_value = schema['exclusiveMinimum'] + 1
  590. if 'maximum' in schema:
  591. max_value = schema['maximum']
  592. elif 'exclusiveMaximum' in schema:
  593. max_value = schema['exclusiveMaximum'] - 1
  594. out = ["("]
  595. _generate_min_max_int(min_value, max_value, out)
  596. out.append(") space")
  597. return self._add_rule(rule_name, ''.join(out))
  598. elif (schema_type == 'object') or (len(schema) == 0):
  599. return self._add_rule(rule_name, self._add_primitive('object', PRIMITIVE_RULES['object']))
  600. else:
  601. assert schema_type in PRIMITIVE_RULES, f'Unrecognized schema: {schema}'
  602. # TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero
  603. return self._add_primitive('root' if rule_name == 'root' else schema_type, PRIMITIVE_RULES[schema_type])
  604. def _add_primitive(self, name: str, rule: BuiltinRule):
  605. n = self._add_rule(name, rule.content)
  606. for dep in rule.deps:
  607. dep_rule = PRIMITIVE_RULES.get(dep) or STRING_FORMAT_RULES.get(dep)
  608. assert dep_rule, f'Rule {dep} not known'
  609. if dep not in self._rules:
  610. self._add_primitive(dep, dep_rule)
  611. return n
  612. def _build_object_rule(self, properties: List[Tuple[str, Any]], required: Set[str], name: str, additional_properties: Optional[Union[bool, Any]]):
  613. prop_order = self._prop_order
  614. # sort by position in prop_order (if specified) then by original order
  615. sorted_props = [kv[0] for _, kv in sorted(enumerate(properties), key=lambda ikv: (prop_order.get(ikv[1][0], len(prop_order)), ikv[0]))]
  616. prop_kv_rule_names = {}
  617. for prop_name, prop_schema in properties:
  618. prop_rule_name = self.visit(prop_schema, f'{name}{"-" if name else ""}{prop_name}')
  619. prop_kv_rule_names[prop_name] = self._add_rule(
  620. f'{name}{"-" if name else ""}{prop_name}-kv',
  621. fr'{self._format_literal(json.dumps(prop_name))} space ":" space {prop_rule_name}'
  622. )
  623. required_props = [k for k in sorted_props if k in required]
  624. optional_props = [k for k in sorted_props if k not in required]
  625. if additional_properties is not None and additional_properties != False:
  626. sub_name = f'{name}{"-" if name else ""}additional'
  627. value_rule = self.visit(additional_properties, f'{sub_name}-value') if isinstance(additional_properties, dict) else \
  628. self._add_primitive('value', PRIMITIVE_RULES['value'])
  629. key_rule = self._add_primitive('string', PRIMITIVE_RULES['string']) if not sorted_props \
  630. else self._add_rule(f'{sub_name}-k', self._not_strings(sorted_props))
  631. prop_kv_rule_names["*"] = self._add_rule(
  632. f'{sub_name}-kv',
  633. f'{key_rule} ":" space {value_rule}'
  634. )
  635. optional_props.append("*")
  636. rule = '"{" space '
  637. rule += ' "," space '.join(prop_kv_rule_names[k] for k in required_props)
  638. if optional_props:
  639. rule += ' ('
  640. if required_props:
  641. rule += ' "," space ( '
  642. def get_recursive_refs(ks, first_is_optional):
  643. [k, *rest] = ks
  644. kv_rule_name = prop_kv_rule_names[k]
  645. comma_ref = f'( "," space {kv_rule_name} )'
  646. if first_is_optional:
  647. res = comma_ref + ('*' if k == '*' else '?')
  648. else:
  649. res = kv_rule_name + (' ' + comma_ref + "*" if k == '*' else '')
  650. if len(rest) > 0:
  651. res += ' ' + self._add_rule(
  652. f'{name}{"-" if name else ""}{k}-rest',
  653. get_recursive_refs(rest, first_is_optional=True)
  654. )
  655. return res
  656. rule += ' | '.join(
  657. get_recursive_refs(optional_props[i:], first_is_optional=False)
  658. for i in range(len(optional_props))
  659. )
  660. if required_props:
  661. rule += ' )'
  662. rule += ' )?'
  663. rule += ' "}" space'
  664. return rule
  665. def format_grammar(self):
  666. return '\n'.join(
  667. f'{name} ::= {rule}'
  668. for name, rule in sorted(self._rules.items(), key=lambda kv: kv[0])
  669. )
  670. def main(args_in = None):
  671. parser = argparse.ArgumentParser(
  672. description='''
  673. Generates a grammar (suitable for use in ./llama-cli) that produces JSON conforming to a
  674. given JSON schema. Only a subset of JSON schema features are supported; more may be
  675. added in the future.
  676. ''',
  677. )
  678. parser.add_argument(
  679. '--prop-order',
  680. default=[],
  681. type=lambda s: s.split(','),
  682. help='''
  683. comma-separated property names defining the order of precedence for object properties;
  684. properties not specified here are given lower precedence than those that are, and
  685. are kept in their original order from the schema. Required properties are always
  686. given precedence over optional properties.
  687. '''
  688. )
  689. parser.add_argument(
  690. '--allow-fetch',
  691. action='store_true',
  692. default=False,
  693. help='Whether to allow fetching referenced schemas over HTTPS')
  694. parser.add_argument(
  695. '--dotall',
  696. action='store_true',
  697. default=False,
  698. help='Whether to treat dot (".") as matching all chars including line breaks in regular expression patterns')
  699. parser.add_argument(
  700. '--raw-pattern',
  701. action='store_true',
  702. default=False,
  703. help='Treats string patterns as raw patterns w/o quotes (or quote escapes)')
  704. parser.add_argument('schema', help='file containing JSON schema ("-" for stdin)')
  705. args = parser.parse_args(args_in)
  706. if args.schema.startswith('https://'):
  707. url = args.schema
  708. import requests
  709. schema = requests.get(url).json()
  710. elif args.schema == '-':
  711. url = 'stdin'
  712. schema = json.load(sys.stdin)
  713. else:
  714. url = f'file://{args.schema}'
  715. with open(args.schema) as f:
  716. schema = json.load(f)
  717. converter = SchemaConverter(
  718. prop_order={name: idx for idx, name in enumerate(args.prop_order)},
  719. allow_fetch=args.allow_fetch,
  720. dotall=args.dotall,
  721. raw_pattern=args.raw_pattern)
  722. schema = converter.resolve_refs(schema, url)
  723. converter.visit(schema, '')
  724. print(converter.format_grammar())
  725. if __name__ == '__main__':
  726. main()