husky/scripts/xq.py

66 lines
2 KiB
Python
Raw Normal View History

2021-08-24 09:09:11 +00:00
#!/usr/bin/env python3
# Based on jeffbr's xq
# a1ba: added unicode output
import argparse
import sys
from typing import Union
from lxml import etree
from lxml.builder import E
from pygments import highlight
from pygments.formatters.other import NullFormatter
from pygments.formatters.terminal import TerminalFormatter
from pygments.lexers.html import XmlLexer
# from . import NAME, DESCRIPTION, VERSION
NAME = 'xq'
VERSION = '0.0.4'
DESCRIPTION = 'like jq but for XML and XPath'
def wrap_in_results(elements: [Union[etree._Element, etree._ElementUnicodeResult]]) -> etree._Element:
results = E.results()
for el in elements:
results.append(E.result(el))
return results
def apply_xpath(infile, xpath_query=None, colorize=False):
try:
parsed = etree.parse(infile, etree.XMLParser(remove_blank_text=True))
except etree.XMLSyntaxError:
parsed = etree.parse(infile, etree.HTMLParser(remove_blank_text=True))
if xpath_query:
matches = parsed.xpath(xpath_query)
results = wrap_in_results(matches)
output = etree.tostring(results, pretty_print=True, encoding='unicode')
else:
output = etree.tostring(parsed, pretty_print=True, encoding='unicode')
formatter = TerminalFormatter() if colorize else NullFormatter()
return highlight(output, XmlLexer(), formatter)
def main():
parser = argparse.ArgumentParser(
prog=NAME,
description=DESCRIPTION,
)
parser.add_argument('-v', '--version', action='version', version='%(prog)s ' + VERSION)
parser.add_argument(
'xpath_query', nargs='?', type=str,
help='XPath query to apply to XML document.'
)
parser.add_argument(
'file', nargs='?', type=argparse.FileType('r'), default=sys.stdin,
help='XML file to process. Defaults to STDIN.',
)
args = parser.parse_args()
sys.stdout.write(apply_xpath(args.file, args.xpath_query, sys.stdout.isatty()))
if __name__ == '__main__':
main()