From 207b970d9412f60b67328ffb4f0a8b23643d50ae Mon Sep 17 00:00:00 2001
From: Batuhan Taskaya <isidentical@gmail.com>
Date: Wed, 8 Dec 2021 18:49:12 +0300
Subject: [PATCH] Automatically enable --stream on server sent events (#1226)

* Automatically enable --stream when used chunked encoding

* try fix 3.6 mock issue

* Only enable on text/event-stream

Co-authored-by: Jakub Roztocil <jakub@roztocil.co>
---
 CHANGELOG.md               |  2 ++
 httpie/output/writer.py    | 13 +++++++++++--
 tests/test_stream.py       | 26 ++++++++++++++++++++++++++
 tests/utils/http_server.py | 16 ++++++++++++++++
 4 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7cdcbb19..5f040ab1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,8 +13,10 @@ This project adheres to [Semantic Versioning](https://semver.org/).
 - Added support for _sending_ multiple HTTP header lines with the same name. ([#130](https://github.com/httpie/httpie/issues/130))
 - Added support for _receiving_ multiple HTTP headers lines with the same name. ([#1207](https://github.com/httpie/httpie/issues/1207))
 - Added support for basic JSON types on `--form`/`--multipart` when using JSON only operators (`:=`/`:=@`). ([#1212](https://github.com/httpie/httpie/issues/1212))
+- Added support for automatically enabling `--stream` when `Content-Type` is `text/event-stream`. ([#376](https://github.com/httpie/httpie/issues/376))
 - Broken plugins will no longer crash the whole application. ([#1204](https://github.com/httpie/httpie/issues/1204))
 
+
 ## [2.6.0](https://github.com/httpie/httpie/compare/2.5.0...2.6.0) (2021-10-14)
 
 [What’s new in HTTPie 2.6.0 →](https://httpie.io/blog/httpie-2.6.0)
diff --git a/httpie/output/writer.py b/httpie/output/writer.py
index 4650264d..cd3eec97 100644
--- a/httpie/output/writer.py
+++ b/httpie/output/writer.py
@@ -2,6 +2,7 @@ import argparse
 import errno
 from typing import IO, TextIO, Tuple, Type, Union
 
+from ..cli.dicts import HTTPHeadersDict
 from ..context import Environment
 from ..models import (
     HTTPRequest,
@@ -110,6 +111,7 @@ def build_output_stream_for_message(
         env=env,
         args=args,
         message_type=message_type,
+        headers=requests_message.headers
     )
     yield from stream_class(
         msg=message_type(requests_message),
@@ -128,16 +130,23 @@ def get_stream_type_and_kwargs(
     env: Environment,
     args: argparse.Namespace,
     message_type: Type[HTTPMessage],
+    headers: HTTPHeadersDict,
 ) -> Tuple[Type['BaseStream'], dict]:
     """Pick the right stream type and kwargs for it based on `env` and `args`.
 
     """
+    is_stream = args.stream
+    if not is_stream and message_type is HTTPResponse:
+        # If this is a response, then check the headers for determining
+        # auto-streaming.
+        is_stream = headers.get('Content-Type') == 'text/event-stream'
+
     if not env.stdout_isatty and not args.prettify:
         stream_class = RawStream
         stream_kwargs = {
             'chunk_size': (
                 RawStream.CHUNK_SIZE_BY_LINE
-                if args.stream
+                if is_stream
                 else RawStream.CHUNK_SIZE
             )
         }
@@ -152,7 +161,7 @@ def get_stream_type_and_kwargs(
                 'encoding_overwrite': args.response_charset,
             })
         if args.prettify:
-            stream_class = PrettyStream if args.stream else BufferedPrettyStream
+            stream_class = PrettyStream if is_stream else BufferedPrettyStream
             stream_kwargs.update({
                 'conversion': Conversion(),
                 'formatting': Formatting(
diff --git a/tests/test_stream.py b/tests/test_stream.py
index 002b245a..55e000a6 100644
--- a/tests/test_stream.py
+++ b/tests/test_stream.py
@@ -2,6 +2,7 @@ import json
 
 import pytest
 import responses
+from unittest.mock import Mock
 
 from httpie.compat import is_windows
 from httpie.cli.constants import PRETTY_MAP
@@ -107,3 +108,28 @@ def test_redirected_stream(httpbin):
     r = http('--pretty=none', '--stream', '--verbose', 'GET',
              httpbin.url + '/get', env=env)
     assert BIN_FILE_CONTENT in r
+
+
+# /drip endpoint produces 3 individual lines,
+# if we set text/event-stream HTTPie should stream
+# it by default. Otherwise, it will buffer and then
+# print.
+@pytest.mark.parametrize('extras, expected', [
+    (
+        ['Accept:text/event-stream'],
+        3
+    ),
+    (
+        ['Accept:text/plain'],
+        1
+    )
+])
+def test_auto_streaming(http_server, extras, expected):
+    env = MockEnvironment()
+    env.stdout.write = Mock()
+    http(http_server + '/drip', *extras, env=env)
+    assert len([
+        call_arg
+        for call_arg in env.stdout.write.call_args_list
+        if b'test' in call_arg[0][0]
+    ]) == expected
diff --git a/tests/utils/http_server.py b/tests/utils/http_server.py
index 42be768b..f09e06c2 100644
--- a/tests/utils/http_server.py
+++ b/tests/utils/http_server.py
@@ -36,6 +36,22 @@ def get_headers(handler):
     handler.end_headers()
 
 
+@TestHandler.handler('GET', '/drip')
+def chunked_drip(handler):
+    handler.send_response(200)
+    accept = handler.headers.get('Accept')
+    if accept is not None:
+        handler.send_header('Content-Type', accept)
+    handler.send_header('Transfer-Encoding', 'chunked')
+    handler.end_headers()
+
+    for _ in range(3):
+        body = 'test\n'
+        handler.wfile.write(f'{len(body):X}\r\n{body}\r\n'.encode('utf-8'))
+
+    handler.wfile.write('0\r\n\r\n'.encode('utf-8'))
+
+
 @pytest.fixture(scope="function")
 def http_server():
     """A custom HTTP server implementation for our tests, that is