added support to use other pygments styles, falback to solarized

This commit is contained in:
Isman Firmansyah 2012-02-29 02:06:36 +07:00
parent 1a88ae647e
commit 553941c98d
2 changed files with 18 additions and 11 deletions

View File

@ -66,6 +66,8 @@ group_only.add_argument('--headers', '-t', dest='print_body',
group_only.add_argument('--body', '-b', dest='print_headers', group_only.add_argument('--body', '-b', dest='print_headers',
action='store_false', default=True, action='store_false', default=True,
help='Print only the response body.') help='Print only the response body.')
parser.add_argument('--style', '-s', dest='style', default='solarized',
help='Output coloring style')
# ``requests.request`` keyword arguments. # ``requests.request`` keyword arguments.
@ -167,11 +169,11 @@ def main():
if args.prettify and sys.stdout.isatty(): if args.prettify and sys.stdout.isatty():
if args.print_headers: if args.print_headers:
status_line = pretty.prettify_http(status_line).strip() status_line = pretty.prettify_http(status_line, args.style).strip()
headers = pretty.prettify_http(headers) headers = pretty.prettify_http(headers, args.style)
if args.print_body: if args.print_body:
body = pretty.prettify_body(body, body = pretty.prettify_body(body,
response.headers['content-type']) response.headers['content-type'], args.style)
if args.print_headers: if args.print_headers:
print status_line print status_line

View File

@ -5,6 +5,7 @@ from pygments.lexers import get_lexer_for_mimetype
from pygments.formatters.terminal256 import Terminal256Formatter from pygments.formatters.terminal256 import Terminal256Formatter
from pygments.lexer import RegexLexer, bygroups from pygments.lexer import RegexLexer, bygroups
from pygments import token from pygments import token
from pygments.styles import get_style_by_name
from . import solarized from . import solarized
@ -24,17 +25,21 @@ class HTTPLexer(RegexLexer):
]} ]}
highlight = partial(pygments.highlight, def _get_highlighter(style_name='solarized'):
formatter=Terminal256Formatter( try:
style=solarized.SolarizedStyle)) style = get_style_by_name(style_name)
highlight_http = partial(highlight, lexer=HTTPLexer()) except pygments.util.ClassNotFound:
style = solarized.SolarizedStyle
return partial(pygments.highlight,
formatter=Terminal256Formatter(style=style))
def prettify_http(headers): def prettify_http(headers, style_name='solarized'):
return highlight_http(headers) highlight = partial(_get_highlighter(style_name), lexer=HTTPLexer())
return highlight(headers)
def prettify_body(content, content_type): def prettify_body(content, content_type, style_name='solarized'):
content_type = content_type.split(';')[0] content_type = content_type.split(';')[0]
if 'json' in content_type: if 'json' in content_type:
content_type = TYPE_JS content_type = TYPE_JS
@ -46,10 +51,10 @@ def prettify_body(content, content_type):
pass pass
try: try:
lexer = get_lexer_for_mimetype(content_type) lexer = get_lexer_for_mimetype(content_type)
highlight = _get_highlighter(style_name)
content = highlight(code=content, lexer=lexer) content = highlight(code=content, lexer=lexer)
if content: if content:
content = content[:-1] content = content[:-1]
except Exception: except Exception:
pass pass
return content return content