|
1 | 1 | # This file was auto-generated by Fern from our API Definition. |
2 | 2 |
|
3 | | -import codecs |
4 | 3 | import re |
5 | 4 | from contextlib import asynccontextmanager, contextmanager |
6 | | -from typing import Any, AsyncGenerator, AsyncIterator, Iterator |
| 5 | +from typing import Any, AsyncGenerator, AsyncIterator, Iterator, cast |
7 | 6 |
|
8 | 7 | import httpx |
9 | 8 | from ._decoders import SSEDecoder |
@@ -46,81 +45,46 @@ def _get_charset(self) -> str: |
46 | 45 | def response(self) -> httpx.Response: |
47 | 46 | return self._response |
48 | 47 |
|
49 | | - @staticmethod |
50 | | - def _normalize_sse_line_endings(buf: str) -> str: |
51 | | - """Normalize line endings per the SSE spec (\\r\\n → \\n, bare \\r → \\n). |
52 | | -
|
53 | | - A trailing \\r is preserved because it may pair with a leading \\n in |
54 | | - the next chunk to form a single \\r\\n terminator. |
55 | | - """ |
56 | | - buf = buf.replace("\r\n", "\n") |
57 | | - if buf.endswith("\r"): |
58 | | - return buf[:-1].replace("\r", "\n") + "\r" |
59 | | - return buf.replace("\r", "\n") |
60 | | - |
61 | 48 | def iter_sse(self) -> Iterator[ServerSentEvent]: |
62 | 49 | self._check_content_type() |
63 | 50 | decoder = SSEDecoder() |
64 | 51 | charset = self._get_charset() |
65 | | - text_decoder = codecs.getincrementaldecoder(charset)(errors="replace") |
66 | 52 |
|
67 | | - buf = "" |
| 53 | + buffer = "" |
68 | 54 | for chunk in self._response.iter_bytes(): |
69 | | - buf += text_decoder.decode(chunk) |
70 | | - buf = self._normalize_sse_line_endings(buf) |
71 | | - |
72 | | - while "\n" in buf: |
73 | | - line, buf = buf.split("\n", 1) |
| 55 | + # Decode chunk using detected charset |
| 56 | + text_chunk = chunk.decode(charset, errors="replace") |
| 57 | + buffer += text_chunk |
| 58 | + |
| 59 | + # Process complete lines |
| 60 | + while "\n" in buffer: |
| 61 | + line, buffer = buffer.split("\n", 1) |
| 62 | + line = line.rstrip("\r") |
74 | 63 | sse = decoder.decode(line) |
| 64 | + # when we reach a "\n\n" => line = '' |
| 65 | + # => decoder will attempt to return an SSE Event |
75 | 66 | if sse is not None: |
76 | 67 | yield sse |
77 | 68 |
|
78 | | - # Flush any remaining bytes from the incremental decoder |
79 | | - buf += text_decoder.decode(b"", final=True) |
80 | | - buf = buf.replace("\r\n", "\n").replace("\r", "\n") |
81 | | - |
82 | | - while "\n" in buf: |
83 | | - line, buf = buf.split("\n", 1) |
| 69 | + # Process any remaining data in buffer |
| 70 | + if buffer.strip(): |
| 71 | + line = buffer.rstrip("\r") |
84 | 72 | sse = decoder.decode(line) |
85 | 73 | if sse is not None: |
86 | 74 | yield sse |
87 | 75 |
|
88 | | - if buf.strip(): |
89 | | - sse = decoder.decode(buf) |
90 | | - if sse is not None: |
91 | | - yield sse |
92 | | - |
93 | 76 | async def aiter_sse(self) -> AsyncGenerator[ServerSentEvent, None]: |
94 | 77 | self._check_content_type() |
95 | 78 | decoder = SSEDecoder() |
96 | | - charset = self._get_charset() |
97 | | - text_decoder = codecs.getincrementaldecoder(charset)(errors="replace") |
98 | | - |
99 | | - buf = "" |
100 | | - async for chunk in self._response.aiter_bytes(): |
101 | | - buf += text_decoder.decode(chunk) |
102 | | - buf = self._normalize_sse_line_endings(buf) |
103 | | - |
104 | | - while "\n" in buf: |
105 | | - line, buf = buf.split("\n", 1) |
| 79 | + lines = cast(AsyncGenerator[str, None], self._response.aiter_lines()) |
| 80 | + try: |
| 81 | + async for line in lines: |
| 82 | + line = line.rstrip("\n") |
106 | 83 | sse = decoder.decode(line) |
107 | 84 | if sse is not None: |
108 | 85 | yield sse |
109 | | - |
110 | | - # Flush any remaining bytes from the incremental decoder |
111 | | - buf += text_decoder.decode(b"", final=True) |
112 | | - buf = buf.replace("\r\n", "\n").replace("\r", "\n") |
113 | | - |
114 | | - while "\n" in buf: |
115 | | - line, buf = buf.split("\n", 1) |
116 | | - sse = decoder.decode(line) |
117 | | - if sse is not None: |
118 | | - yield sse |
119 | | - |
120 | | - if buf.strip(): |
121 | | - sse = decoder.decode(buf) |
122 | | - if sse is not None: |
123 | | - yield sse |
| 86 | + finally: |
| 87 | + await lines.aclose() |
124 | 88 |
|
125 | 89 |
|
126 | 90 | @contextmanager |
|
0 commit comments