github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/aistore/sdk/object_reader.py (about)

     1  from typing import Iterator
     2  
     3  import requests
     4  from requests.structures import CaseInsensitiveDict
     5  
     6  from aistore.sdk.const import DEFAULT_CHUNK_SIZE
     7  from aistore.sdk.object_attributes import ObjectAttributes
     8  
     9  
    10  class ObjectReader:
    11      """
    12      Represents the data returned by the API when getting an object, including access to the content stream and object
    13      attributes
    14      """
    15  
    16      def __init__(
    17          self,
    18          response_headers: CaseInsensitiveDict,
    19          stream: requests.Response,
    20          chunk_size: int = DEFAULT_CHUNK_SIZE,
    21      ):
    22          self._chunk_size = chunk_size
    23          self._stream = stream
    24          self._attributes = ObjectAttributes(response_headers)
    25  
    26      @property
    27      def attributes(self) -> ObjectAttributes:
    28          """
    29          Object metadata attributes
    30  
    31          Returns:
    32              Object attributes parsed from the headers returned by AIS
    33          """
    34          return self._attributes
    35  
    36      def read_all(self) -> bytes:
    37          """
    38          Read all byte data from the object content stream.
    39          This uses a bytes cast which makes it slightly slower and requires all object content to fit in memory at once
    40          Returns:
    41              Object content as bytes
    42  
    43          """
    44          obj_arr = bytearray()
    45          for chunk in self:
    46              obj_arr.extend(chunk)
    47          return bytes(obj_arr)
    48  
    49      def raw(self) -> bytes:
    50          """
    51          Returns: Raw byte stream of object content
    52          """
    53          return self._stream.raw
    54  
    55      def __iter__(self) -> Iterator[bytes]:
    56          """
    57          Creates a generator to read the stream content in chunks
    58          Returns:
    59              An iterator with access to the next chunk of bytes
    60          """
    61          try:
    62              for chunk in self._stream.iter_content(chunk_size=self._chunk_size):
    63                  yield chunk
    64          finally:
    65              self._stream.close()