|
| 1 | +import sys |
| 2 | +from pathlib import Path |
| 3 | + |
| 4 | +from tree_sitter_markdown import language, inline_language |
| 5 | +from tree_sitter import Parser, Language |
| 6 | + |
| 7 | +block_language = Language(language()) |
| 8 | +block_parser = Parser(block_language) |
| 9 | + |
| 10 | +link_reference_def = block_language.query("(link_reference_definition) @node") |
| 11 | +section = block_language.query("(section) @node") |
| 12 | +inline = block_language.query("(inline) @node") |
| 13 | + |
| 14 | +inline_language = Language(inline_language()) |
| 15 | +inline_parser = Parser(inline_language) |
| 16 | +inline_link = inline_language.query("(inline_link) @node") |
| 17 | + |
| 18 | +def node_matches(query, node): |
| 19 | + for idx, match in query.matches(node): |
| 20 | + yield match["node"][0] |
| 21 | + |
| 22 | +def child_for_type(node, typ): |
| 23 | + return [c for c in node.children if c.type == typ][0] |
| 24 | + |
| 25 | +def main(filenames: list[str]) -> int: |
| 26 | + exit_code = 0 |
| 27 | + for f in filenames: |
| 28 | + edits = [] |
| 29 | + link_references = {} |
| 30 | + lines_to_add = [] |
| 31 | + |
| 32 | + buf = Path(f).read_bytes() |
| 33 | + tree = block_parser.parse(buf) |
| 34 | + # .root_node = document |
| 35 | + # .children[] = section |
| 36 | + for node in node_matches(link_reference_def, tree.root_node): |
| 37 | + link_references[child_for_type(node, "link_label").text] = child_for_type(node, "link_destination").text |
| 38 | + |
| 39 | + for inline_node in node_matches(inline, tree.root_node): |
| 40 | + inline_tree = inline_parser.parse(inline_node.text) |
| 41 | + |
| 42 | + for link in node_matches(inline_link, inline_tree.root_node): |
| 43 | + dest = child_for_type(link, "link_destination") |
| 44 | + if len(dest.text) > 30: |
| 45 | + link_text = b"[" + child_for_type(link, "link_text").text + b"]" |
| 46 | + if link_text not in link_references: |
| 47 | + link_references[link_text] = dest.text |
| 48 | + # TODO link_text might have newlines, should replace with single space? |
| 49 | + lines_to_add.append(link_text + b": " + dest.text) |
| 50 | + edits.append((link.start_byte + inline_node.start_byte, link.end_byte + inline_node.start_byte, link_text + b"[]")) |
| 51 | + |
| 52 | + if edits: |
| 53 | + for i, j, new_bytes in sorted(edits, reverse=True): |
| 54 | + buf = buf[:i] + new_bytes + buf[j:] |
| 55 | + |
| 56 | + while not buf.endswith(b"\n\n"): |
| 57 | + buf += b"\n" |
| 58 | + |
| 59 | + for line_to_add in lines_to_add: |
| 60 | + buf += line_to_add + b"\n" |
| 61 | + |
| 62 | + # print(buf.decode()) |
| 63 | + Path(f).write_bytes(buf) |
| 64 | + |
| 65 | + return exit_code |
| 66 | + |
| 67 | +if __name__ == "__main__": |
| 68 | + sys.exit(main(sys.argv[1:])) |
0 commit comments