Added --style

Closes #6. Thanks, @iromli.
This commit is contained in:
Jakub Roztočil 2012-03-02 01:39:22 +01:00
commit 94c605fac1
3 changed files with 45 additions and 37 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

@ -1,16 +1,23 @@
import json import re
from functools import partial
import pygments
import os import os
import json
import pygments
from pygments import token
from pygments.util import ClassNotFound
from pygments.lexers import get_lexer_for_mimetype from pygments.lexers import get_lexer_for_mimetype
from pygments.formatters.terminal256 import Terminal256Formatter from pygments.formatters.terminal256 import Terminal256Formatter
from pygments.formatters.terminal import TerminalFormatter from pygments.formatters.terminal import TerminalFormatter
from pygments.lexer import RegexLexer, bygroups from pygments.lexer import RegexLexer, bygroups
from pygments import token from pygments.styles import get_style_by_name, STYLE_MAP
from . import solarized from . import solarized
DEFAULT_STYLE = 'solarized'
AVAILABLE_STYLES = [DEFAULT_STYLE] + STYLE_MAP.keys()
TYPE_JS = 'application/javascript' TYPE_JS = 'application/javascript'
FORMATTER = (Terminal256Formatter
if os.environ.get('TERM', None) == 'xterm-256color'
else TerminalFormatter)
class HTTPLexer(RegexLexer): class HTTPLexer(RegexLexer):
@ -25,23 +32,20 @@ class HTTPLexer(RegexLexer):
(r'(.*?:)(.+)', bygroups(token.Name, token.String)) (r'(.*?:)(.+)', bygroups(token.Name, token.String))
]} ]}
if os.environ['TERM'] == 'xterm-256color':
formatter = Terminal256Formatter
else:
formatter = TerminalFormatter
class PrettyHttp(object):
highlight = partial(pygments.highlight, def __init__(self, style_name):
formatter=formatter( if style_name == 'solarized':
style=solarized.SolarizedStyle)) style = solarized.SolarizedStyle
highlight_http = partial(highlight, lexer=HTTPLexer()) else:
style = get_style_by_name(style_name)
self.formatter = FORMATTER(style=style)
def headers(self, content):
return pygments.highlight(content, HTTPLexer(), self.formatter).strip()
def prettify_http(headers): def body(self, content, content_type):
return highlight_http(headers)
def prettify_body(content, content_type):
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
@ -53,10 +57,9 @@ def prettify_body(content, content_type):
pass pass
try: try:
lexer = get_lexer_for_mimetype(content_type) lexer = get_lexer_for_mimetype(content_type)
content = highlight(code=content, lexer=lexer) except ClassNotFound:
if content:
content = content[:-1]
except Exception:
pass
return content return content
content = pygments.highlight(content, lexer, self.formatter)
# TODO: Make sure the leading/trailing whitespaces remain.
# Some of the Pygments styles add superfluous line breaks.
return content.strip()

3
tests.py Normal file
View File

@ -0,0 +1,3 @@
import unittest