Copyright (C) 2025 HERE Global B.V. and its affiliate(s).

All rights reserved.

This software and other materials contain proprietary information

controlled by HERE and are protected by applicable copyright legislation.

Any use and utilization of this software and other materials and

disclosure to any third parties is conditional upon having a separate

agreement with HERE for the access, use, utilization or disclosure of this

software. In the absence of such agreement, the use of the software is not

allowed.

"""
This module contains classes to implement streaming requests with retries.
"""

import os
import time
from collections.abc import Iterator
from io import BytesIO, RawIOBase
from typing import Optional, Union, cast

import backoff
import json_stream
from here.platform.api.base_api import BaseApi
from json_stream.base import (
StreamingJSONList,
StreamingJSONObject,
TransientAccessException,
TransientStreamingJSONBase,
)
from requests import Response
from requests.exceptions import ChunkedEncodingError, ConnectionError

[docs]
class ChunkedGet(RawIOBase, Iterator[bytes]):
"""
Custom iterator for iterating content from a GET request that allows retries on error by using
the Range HTTP header to restart where it left off.

This may also be used as a file-like stream.
"""

def init(self, api: BaseApi, response: Response, chunk_size: int):
"""
Initializes this with the response to stream.

:param api: the API the request was made with.
:param response: the response to the request.
:param chunk_size: the desired amount of data to request at a time.
"""
self._api = api
self._response = response
self._chunk_size = chunk_size
self._iterator: Iterator[bytes] = response.iter_content(chunk_size)
self._received_bytes = 0
self._cur_chunk: Optional[bytes] = None
self._cur_offset = 0
self._retry = backoff.expo()
self._wait_time = 0.0

[docs]
def readable(self) -> bool:
"""
Denotes to the RawIO API that this stream is readable.

:return: True to indicate the stream is readable.
"""
return True

[docs]
def read(self, size: int = -1) -> bytes:
"""
Reads bytes from the stream.

:param size: the requested number of bytes to read. Fewer bytes may be read if it would
require multiple requests. If negative, all remaining data will be read as if readall()
was called.
:return: the read bytes.
"""
if size < 0:
return self.readall()

First read from any remaining data from the previous chunk.

if self._cur_chunk:
remaining = len(self._cur_chunk) - self._cur_offset
read_size = min(remaining, size)
res = self._cur_chunk[self._cur_offset : self._cur_offset + read_size]
if read_size == remaining:
self._cur_chunk = None
self._cur_offset = 0
else:
self._cur_offset += read_size
return res

Then read from the next chunk.

try:
chunk = self._next_chunk()
if len(chunk) <= size:
return chunk
else:
res = chunk[:size]
self._cur_chunk = chunk
self._cur_offset = size
return res
except StopIteration:
return b""

[docs]
def readall(self) -> bytes:
"""
Reads all remaining bytes from the stream.

:return: the read bytes.
"""
with BytesIO() as temp_buffer:

First read from any remaining data from the previous chunk.

if self._cur_chunk:
temp_buffer.write(memoryview(self._cur_chunk)[self._cur_offset :])
self._cur_chunk = None
self._cur_offset = 0

Then read from following chunks.

while True:
try:
temp_buffer.write(self._next_chunk())
except StopIteration:
break

return temp_buffer.getvalue()

[docs]
def readinto(self, b) -> int:
"""
Reads bytes into a pre-allocated buffer.

:param b: a buffer-like object to read into.
:return: the number of bytes read.
"""
size = len(b)
offset = 0

First read from any remaining data from the previous chunk.

if self._cur_chunk:
remaining = len(self._cur_chunk) - self._cur_offset
read_size = min(remaining, size)
b[offset : offset + read_size] = memoryview(self._cur_chunk)[
self._cur_offset : self._cur_offset + read_size
]
offset += read_size
size -= read_size
if read_size == remaining:
self._cur_chunk = None
self._cur_offset = 0
else:
assert size == 0
self._cur_offset += read_size

Then read from following chunks.

while size > 0:
try:
this_chunk = self._next_chunk()
chunk_size = len(this_chunk)
if chunk_size <= size:
b[offset : offset + chunk_size] = this_chunk
offset += chunk_size
size -= chunk_size
else:
b[offset : offset + size] = memoryview(this_chunk)[:size]
self._cur_chunk = this_chunk
self._cur_offset = size
offset += size
size = 0
except StopIteration:
break

return offset

[docs]
def seekable(self) -> bool:
"""
Allow seeking to skip forward, allowing for buffered reading within json_stream.

:return: True to state that this is seekable.
"""
return True

[docs]
def seek(self, offset: int, whence=os.SEEK_SET) -> int:
"""
Seeks in the stream. Only seeking forward is supported.

:param offset: the offset to seek by.
:param whence: the position offset is relative to.
:return: the new position in the stream.
:raises OSError: if seeking backward.
"""
if whence == os.SEEK_SET:
offset -= self.tell()
elif whence != os.SEEK_CUR:
raise OSError("Seeking may only be forward for ChunkedGet.")

if offset < 0:
raise OSError("Seeking may only be forward for ChunkedGet.")

First skip from any remaining data from the previous chunk.

if self._cur_chunk:
remaining = len(self._cur_chunk) - self._cur_offset
read_size = min(remaining, offset)
offset -= read_size
if read_size == remaining:
self._cur_chunk = None
self._cur_offset = 0
else:
assert offset == 0
self._cur_offset += read_size

Then skip from following chunks.

while offset > 0:
try:
this_chunk = self._next_chunk()
if len(this_chunk) <= offset:
offset -= len(this_chunk)
else:
self._cur_chunk = this_chunk
self._cur_offset = offset
offset = 0
except StopIteration:
break
return self.tell()

[docs]
def tell(self) -> int:
"""
Tells the current position in the stream.

:return: the current position in bytes.
"""
position = self._received_bytes
if self._cur_chunk:
position -= len(self._cur_chunk) - self._cur_offset
return position

[docs]
def close(self):
"""
Closes the stream and releases any resources.
"""
self._response.close()

def iter(self) -> Iterator[bytes]:
"""
Gets the iterator to iterate over the bytes.

:return: self as the iterator.
"""
return self

def next(self) -> bytes:
"""
Gets the next chunk of bytes for iteration.

:return: the next chunk of bytes.
:rases StopIteration: once iteration has completed.
"""
if self._cur_chunk:
res = self._cur_chunk[self._cur_offset :]
self._cur_offset = 0
self._cur_chunk = None
return res
return self._next_chunk()

def _next_chunk(self) -> bytes:
while True:
try:
chunk = next(self._iterator)
self._received_bytes += len(chunk)

Reset retries once we've had a successful request.

self._retry = backoff.expo()
self._wait_time = 0.0
return chunk
except (ChunkedEncodingError, ConnectionError):
if self._wait_time >= self._api.application_config.retry_max_time:
raise

next_wait = next(self._retry)
if next_wait:
self._wait_time += next_wait
time.sleep(next_wait)

request = self._response.request
assert request.url is not None
assert request.headers is not None
headers = request.headers.copy()
headers["Range"] = f"bytes={self._received_bytes}-"
self._response = self._api.get(url=request.url, headers=headers, stream=True)
self._iterator = self._response.iter_content(self._chunk_size)

Loop back around to try again.

[docs]
def stream_json_response(api: BaseApi, response: Response, chunk_size: int) -> StreamingJSONObject:
"""
Streams loading a JSON response to avoid keeping the full document in memory.

:param api: the API the request was made with.
:param response: the response to the request.
:param chunk_size: the desired amount of data to request at a time.
:return: an object to access the streamed JSON data.
"""
return cast(StreamingJSONObject, json_stream.load(ChunkedGet(api, response, chunk_size)))

[docs]
def find_list_of_objects(
parent: Union[dict, StreamingJSONObject], field: str
) -> Union[list[dict], StreamingJSONList]:
"""
Finds a list of objects to iterate over from either a fully loaded JSON dict or streaming JSON.
When performing streaming parsing, this will persist the individual objects so lookups can
succeed regardless of the order they are present in the JSON document. This assumes a
potentially enormous list of small objects.

:param parent: the parent object to look in.
:param field: the field to look up.
:return: the list of objects. If field is not present, an empty list will be returned.
"""
try:
member = parent.get(field, [])
if isinstance(member, TransientStreamingJSONBase):
return cast(StreamingJSONList, member.persistent())
return member
except TransientAccessException:
return []