diff --git a/docs/README.md b/docs/README.md index 51dff51aec..f8b0eb0410 100644 --- a/docs/README.md +++ b/docs/README.md @@ -2068,7 +2068,7 @@ You can use the `--stream, -S` flag to make two things happen: 1. The output is flushed in much smaller chunks without any buffering, which makes HTTPie behave kind of like `tail -f` for URLs. 2. Streaming becomes enabled even when the output is prettified: It will be applied to each line of the response and flushed immediately. This makes it possible to have a nice output for long-lived requests, such as one to the [Twitter streaming API](https://developer.twitter.com/en/docs/tutorials/consuming-streaming-data). -The `--stream` option is automatically enabled when the response headers include `Content-Type: text/event-stream`. +The `--stream` option is automatically enabled when the response `Content-Type` is `text/event-stream`, `application/x-ndjson`, or `application/ndjson`. ### Example use cases diff --git a/httpie/output/streams.py b/httpie/output/streams.py index 811093808a..dd4640d686 100644 --- a/httpie/output/streams.py +++ b/httpie/output/streams.py @@ -216,12 +216,20 @@ def iter_body(self) -> Iterable[bytes]: yield self.process_body(line) + lf first_chunk = False + SSE_DATA_PREFIX = 'data:' + SSE_MIME = 'text/event-stream' + JSON_MIME = 'application/json' + def process_body(self, chunk: Union[str, bytes]) -> bytes: if not isinstance(chunk, str): # Text when a converter has been used, # otherwise it will always be bytes. chunk = self.decode_chunk(chunk) - chunk = self.formatting.format_body(content=chunk, mime=self.mime) + mime = self.mime + if mime == self.SSE_MIME and chunk.lstrip().startswith(self.SSE_DATA_PREFIX): + # SSE data lines may contain JSON; use JSON mime for highlighting + mime = self.JSON_MIME + chunk = self.formatting.format_body(content=chunk, mime=mime) return smart_encode(chunk, self.output_encoding) diff --git a/httpie/output/writer.py b/httpie/output/writer.py index 4a2949bce2..23345b92e6 100644 --- a/httpie/output/writer.py +++ b/httpie/output/writer.py @@ -23,6 +23,12 @@ MESSAGE_SEPARATOR = '\n\n' MESSAGE_SEPARATOR_BYTES = MESSAGE_SEPARATOR.encode() +AUTO_STREAMING_CONTENT_TYPES = frozenset({ + 'text/event-stream', + 'application/x-ndjson', + 'application/ndjson', +}) + def write_message( requests_message: RequestsMessage, @@ -167,7 +173,7 @@ def get_stream_type_and_kwargs( raw_content_type_header = headers.get('Content-Type', None) if raw_content_type_header: content_type_header, _ = parse_content_type_header(raw_content_type_header) - is_stream = (content_type_header == 'text/event-stream') + is_stream = (content_type_header in AUTO_STREAMING_CONTENT_TYPES) if not env.stdout_isatty and not prettify_groups: stream_class = RawStream diff --git a/tests/test_stream.py b/tests/test_stream.py index b0b9b8bde8..db7c41be59 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -115,10 +115,9 @@ def test_redirected_stream(httpbin): 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. +# /drip endpoint produces 3 individual lines. +# Auto-streaming content types (SSE, NDJSON) should stream +# line by line; other types should buffer into a single write. @pytest.mark.parametrize('extras, expected', [ ( ['Accept:text/event-stream'], @@ -128,6 +127,18 @@ def test_redirected_stream(httpbin): ['Accept:text/event-stream; charset=utf-8'], 3 ), + ( + ['Accept:application/x-ndjson'], + 3 + ), + ( + ['Accept:application/ndjson'], + 3 + ), + ( + ['Accept:application/x-ndjson; charset=utf-8'], + 3 + ), ( ['Accept:text/plain'], 1 @@ -144,6 +155,35 @@ def test_auto_streaming(http_server, extras, expected): ]) == expected +def test_ndjson_auto_streaming(http_server): + """NDJSON responses should auto-stream line by line.""" + env = MockEnvironment() + env.stdout.write = Mock() + http(http_server + '/ndjson', env=env) + json_writes = [ + call_arg + for call_arg in env.stdout.write.call_args_list + if 'message' in call_arg[0][0] + ] + assert len(json_writes) == 3 + + +def test_ndjson_json_formatting(http_server): + """NDJSON lines should be pretty-printed as JSON.""" + r = http('--pretty=format', http_server + '/ndjson') + assert '"message": "test"' in r + assert '"index": 0' in r + + +def test_sse_json_highlighting(http_server): + """SSE data lines containing JSON should get syntax highlighting.""" + r = http('--pretty=colors', http_server + '/sse-json') + # With colors enabled, JSON keys/values should have ANSI escape codes + # The raw 'data:' prefix and JSON content should both appear + assert 'data:' in r + assert 'message' in r + + def test_streaming_encoding_detection(http_server): r = http('--stream', http_server + '/stream/encoding/random') assert ASCII_FILE_CONTENT in r diff --git a/tests/utils/http_server.py b/tests/utils/http_server.py index 86cc069c57..c401c06fc7 100644 --- a/tests/utils/http_server.py +++ b/tests/utils/http_server.py @@ -42,22 +42,53 @@ def get_headers(handler): handler.end_headers() -@TestHandler.handler('GET', '/drip') -def chunked_drip(handler): +def _send_chunked_response(handler, content_type, chunks): + """Send a chunked HTTP response with the given content type and body chunks.""" handler.send_response(200) - accept = handler.headers.get('Accept') - if accept is not None: - handler.send_header('Content-Type', accept) + if content_type is not None: + handler.send_header('Content-Type', content_type) handler.send_header('Transfer-Encoding', 'chunked') handler.end_headers() - for _ in range(3): - body = 'test\n' + for body in chunks: 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')) +@TestHandler.handler('GET', '/drip') +def chunked_drip(handler): + _send_chunked_response( + handler, + content_type=handler.headers.get('Accept'), + chunks=['test\n' for _ in range(3)], + ) + + +@TestHandler.handler('GET', '/ndjson') +def ndjson_stream(handler): + _send_chunked_response( + handler, + content_type='application/x-ndjson', + chunks=[ + json.dumps({"index": i, "message": "test"}) + '\n' + for i in range(3) + ], + ) + + +@TestHandler.handler('GET', '/sse-json') +def sse_json_stream(handler): + _send_chunked_response( + handler, + content_type='text/event-stream', + chunks=[ + f'data: {json.dumps({"index": i, "message": "test"})}\n\n' + for i in range(3) + ], + ) + + @TestHandler.handler('GET', '/stream/encoding/random') def random_encoding(handler): from tests.fixtures import ASCII_FILE_CONTENT, FILE_CONTENT as UNICODE_FILE_CONTENT