From 5653b9c6a010c7e64d778c350644fa722a289139 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Roztoc=CC=8Cil?= <jakub@roztocil.name>
Date: Sun, 26 Feb 2012 01:37:28 +0100
Subject: [PATCH] Added --headers and --body to limit the output.

---
 httpie/httpie.py | 35 +++++++++++++++++++++++++++--------
 httpie/pretty.py | 20 ++++++++++++--------
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/httpie/httpie.py b/httpie/httpie.py
index 9b6c1dbd..cab1ffb5 100755
--- a/httpie/httpie.py
+++ b/httpie/httpie.py
@@ -6,7 +6,7 @@ import argparse
 from collections import namedtuple
 import requests
 from requests.structures import CaseInsensitiveDict
-from .pretty import prettify
+from . import pretty
 
 
 __author__ = 'Jakub Roztocil'
@@ -46,17 +46,26 @@ parser = argparse.ArgumentParser(
 
 
 # Content type.
-group = parser.add_mutually_exclusive_group(required=False)
-group.add_argument('--json', '-j', action='store_true',
+group_type = parser.add_mutually_exclusive_group(required=False)
+group_type.add_argument('--json', '-j', action='store_true',
                    help='Serialize data items as a JSON object and set'
                         ' Content-Type to application/json, if not specified.')
-group.add_argument('--form', '-f', action='store_true',
+group_type.add_argument('--form', '-f', action='store_true',
                    help='Serialize data items as form values and set'
                         ' Content-Type to application/x-www-form-urlencoded,'
                         ' if not specified.')
 
+# Output options.
 parser.add_argument('--ugly', '-u', help='Do not prettify the response.',
                      dest='prettify', action='store_false', default=True)
+group_only = parser.add_mutually_exclusive_group(required=False)
+group_only.add_argument('--headers', '-t', dest='print_body',
+                        action='store_false', default=True,
+                        help='Print only the response headers.')
+group_only.add_argument('--body', '-b', dest='print_headers',
+                        action='store_false', default=True,
+                        help='Print only the response body.')
+
 
 # ``requests.request`` keyword arguments.
 parser.add_argument('--auth', help='username:password',
@@ -136,7 +145,7 @@ def main():
 
     # Display the response.
     original = response.raw._original_response
-    response_bits = (
+    status_line, headers, body = (
         u'HTTP/{version} {status} {reason}'.format(
             version='.'.join(str(original.version)),
             status=original.status, reason=original.reason,
@@ -146,10 +155,20 @@ def main():
     )
 
     if args.prettify and sys.stdout.isatty():
-        response_bits = prettify(response.headers['content-type'], *response_bits)
-
-    print u'\n'.join(response_bits)
+        if args.print_headers:
+            status_line = pretty.prettify_http(status_line)
+            headers = pretty.prettify_http(headers)
+        if args.print_body:
+            body = pretty.prettify_body(body,
+                response.headers['content-type'])
 
+    if args.print_headers:
+        print status_line
+        print headers
+    if args.print_body:
+        if args.print_headers:
+            print
+        print body
 
 if __name__ == '__main__':
     main()
diff --git a/httpie/pretty.py b/httpie/pretty.py
index 6a0d914d..ee6eff4f 100644
--- a/httpie/pretty.py
+++ b/httpie/pretty.py
@@ -28,22 +28,26 @@ highlight = partial(pygments.highlight,
 highlight_http = partial(highlight, lexer=HTTPLexer())
 
 
-def prettify(content_type, status_line, headers, body):
-    content_type = content_type.split(';')[0]
+def prettify_http(headers):
+    return highlight_http(headers)[:-1]
 
+
+def prettify_body(content, content_type):
+    content_type = content_type.split(';')[0]
     if 'json' in content_type:
         content_type = TYPE_JS
         try:
             # Indent JSON
-            body = json.dumps(json.loads(body), sort_keys=True, indent=4)
+            content = json.dumps(json.loads(content),
+                                sort_keys=True, indent=4)
         except Exception:
             pass
-
     try:
-        body = highlight(code=body, lexer=get_lexer_for_mimetype(content_type))
+        lexer = get_lexer_for_mimetype(content_type)
+        content = highlight(code=content, lexer=lexer)
+        if content:
+            content = content[:-1]
     except Exception:
         pass
+    return content
 
-    return (highlight_http(code=status_line).strip(),
-            highlight_http(code=headers),
-            body)