From bced559496de7735b6cc37c77bf00d5c241542c0 Mon Sep 17 00:00:00 2001
From: Jakob Kramer <jakob.kramer@gmx.de>
Date: Sat, 28 Apr 2012 14:13:40 +0200
Subject: [PATCH] use PrettyHttp class; working --headers and --body

---
 httpie/__main__.py    |  8 +++-----
 httpie/httpmessage.py | 24 +++++++++---------------
 httpie/pretty.py      | 37 ++++++++++++++++++++++++++++---------
 3 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/httpie/__main__.py b/httpie/__main__.py
index 3cf60c46..8cc48001 100644
--- a/httpie/__main__.py
+++ b/httpie/__main__.py
@@ -84,14 +84,13 @@ def _get_output(args, stdout_isatty, response):
     do_output_response = (cliparse.OUT_RESP_HEADERS in args.output_options
                           or cliparse.OUT_RESP_BODY in args.output_options)
 
-    formatter = pretty.formatter(args.style) if do_prettify else None
+    prettifier = pretty.PrettyHttp(args.style) if do_prettify else None
     output = []
 
     if do_output_request:
         output.append(httpmessage.format(
             message=httpmessage.from_request(response.request),
-            pretty=do_prettify,
-            formatter=formatter,
+            prettifier=prettifier,
             with_headers=cliparse.OUT_REQ_HEADERS in args.output_options,
             with_body=cliparse.OUT_REQ_BODY in args.output_options
         ))
@@ -101,8 +100,7 @@ def _get_output(args, stdout_isatty, response):
     if do_output_response:
         output.append(httpmessage.format(
             message=httpmessage.from_response(response),
-            pretty=do_prettify,
-            formatter=formatter,
+            prettifier=prettifier,
             with_headers=cliparse.OUT_RESP_HEADERS in args.output_options,
             with_body=cliparse.OUT_RESP_BODY in args.output_options
         ))
diff --git a/httpie/httpmessage.py b/httpie/httpmessage.py
index d063a807..2a9b31c6 100644
--- a/httpie/httpmessage.py
+++ b/httpie/httpmessage.py
@@ -1,7 +1,3 @@
-import json
-
-from pygments import highlight
-from pygments.lexers import HttpLexer
 from requests.compat import urlparse
 
 
@@ -48,26 +44,24 @@ def from_response(response):
         content_type=response_headers.get('Content-Type'))
 
 
-def format(message, formatter=None, pretty=False,
+def format(message, prettifier=None,
            with_headers=True, with_body=True):
     """Return a `unicode` representation of `message`. """
+    pretty = prettifier is not None
     bits = []
+
     if with_headers:
         bits.append(message.line)
         bits.append(message.headers)
+        if pretty:
+            bits = [prettifier.headers('\n'.join(bits))]
         if with_body and message.body:
             bits.append('\n')
+
     if with_body and message.body:
-        if (pretty and message.content_type and
-                message.content_type.split(';')[0] == 'application/json'):
-            # Indent and sort the JSON data.
-            bits.append(json.dumps(json.loads(message.body),
-                                   sort_keys=True, indent=4))
+        if pretty and message.content_type:
+            bits.append(prettifier.body(message.body, message.content_type))
         else:
             bits.append(message.body)
-    bits.append('\n')
-    result = '\n'.join(bit.strip() for bit in bits)
-    if pretty:
-        return highlight(result, HttpLexer(), formatter)
 
-    return result
+    return '\n'.join(bit.strip() for bit in bits)
diff --git a/httpie/pretty.py b/httpie/pretty.py
index cc40b605..deb82bb3 100644
--- a/httpie/pretty.py
+++ b/httpie/pretty.py
@@ -3,12 +3,11 @@ import json
 
 import pygments
 
-from pygments import token
 from pygments.util import ClassNotFound
-from pygments.styles import get_style_by_name, STYLE_MAP
 from pygments.lexers import get_lexer_for_mimetype, HttpLexer
-from pygments.formatters.terminal import TerminalFormatter
 from pygments.formatters.terminal256 import Terminal256Formatter
+from pygments.formatters.terminal import TerminalFormatter
+from pygments.styles import get_style_by_name, STYLE_MAP
 
 from . import solarized
 
@@ -20,11 +19,31 @@ FORMATTER = (Terminal256Formatter
              else TerminalFormatter)
 
 
-def formatter(style_name):
-    if style_name == 'solarized':
-        style = solarized.SolarizedStyle
+class PrettyHttp(object):
 
-    else:
-        style = get_style_by_name(style_name)
+    def __init__(self, style_name):
+        if style_name == 'solarized':
+            style = solarized.SolarizedStyle
+        else:
+            style = get_style_by_name(style_name)
+        self.formatter = FORMATTER(style=style)
 
-    return FORMATTER(style=style)
+    def headers(self, content):
+        return pygments.highlight(content, HttpLexer(), self.formatter)
+
+    def body(self, content, content_type):
+        content_type = content_type.split(';')[0]
+        try:
+            lexer = get_lexer_for_mimetype(content_type)
+        except ClassNotFound:
+            return content
+
+        if content_type == 'application/json':
+            try:
+                # Indent and sort the JSON data.
+                content = json.dumps(json.loads(content),
+                                     sort_keys=True, indent=4)
+            except:
+                pass
+
+        return pygments.highlight(content, lexer, self.formatter)