Merge pull request #346 from pcattori/response-ndjson

response.ndjson utility

CHANGELOG.md

@@ -30,6 +30,8 @@
       - functions: `tc.session.from_auth`
     - `tc.url` module
       - `tc.URL` type
+    - `tc.response` module
+      - functions: `successful`, `ndjson`
 
   **BUG FIXES**
   - Links from our docs to the `requests` docs were outdated. Links have been updated to point to the new `requests` docs URL.

docs/beta.md

@@ -9,4 +9,5 @@
   * [Auth](beta/auth)
   * [Dataset](beta/datasets)
   * [Instance](beta/instance)
+  * [Response](beta/response)
   * [Session](beta/session)

docs/beta/response.rst

@@ -0,0 +1,7 @@
+Response
+========
+
+Utilities for working with :class:`requests.Response` .
+
+.. autofunction:: tamr_client.response.successful
+.. autofunction:: tamr_client.response.ndjson

stubs/responses.pyi

@@ -10,6 +10,7 @@ PUT: str
 def add(
     method: Optional[str] = None,
     url: Optional[str] = None,
+    body: Optional[str] = None,
     status: Optional[int] = None,
     json: Optional[JsonDict] = None,
 ): ...

tamr_client/response.py

@@ -1,7 +1,11 @@
+import json
 import logging
+from typing import Iterator
 
 import requests
 
+from tamr_client.types import JsonDict
+
 logger = logging.getLogger(__name__)
 
 
@@ -25,3 +29,29 @@ def successful(response: requests.Response) -> requests.Response:
         )
         raise e
     return response
+
+
+def ndjson(response: requests.Response, **kwargs) -> Iterator[JsonDict]:
+    """Stream newline-delimited JSON from the response body
+
+    Analog to :func:`requests.Response.json` but for ``.ndjson``-formatted body.
+
+    **Recommended**: For memory efficiency, use ``stream=True`` when sending the request corresponding to this response.
+
+    Args:
+        response: Response whose body should be streamed as newline-delimited JSON.
+        **kwargs: Keyword arguments passed to underlying :func:`requests.Response.iter_lines` call.
+
+    Returns
+        Each line of the response body, parsed as JSON
+
+    Example:
+        >>> import tamr_client as tc
+        >>> s = tc.session.from_auth(...)
+        >>> r = s.get(..., stream=True)
+        >>> for data in tc.response.ndjson(r):
+        ...     assert data['my key'] == 'my_value'
+
+    """
+    for line in response.iter_lines(**kwargs):
+        yield json.loads(line)

tests/tamr_client/test_response.py

@@ -0,0 +1,24 @@
+import json
+
+import responses
+
+import tamr_client as tc
+import tests.tamr_client.utils as utils
+
+
+@responses.activate
+def test_ndjson():
+    s = utils.session()
+
+    records = [{"a": 1}, {"b": 2}, {"c": 3}]
+    url = tc.URL(path="datasets/1/records")
+    responses.add(
+        responses.GET, str(url), body="\n".join(json.dumps(x) for x in records)
+    )
+
+    r = s.get(str(url))
+
+    ndjson = list(tc.response.ndjson(r))
+    assert len(ndjson) == 3
+    for record in ndjson:
+        assert record in records