From 198935a306adff45adcbdc298a2bcf8f4383ff2e Mon Sep 17 00:00:00 2001
From: Abhishek Gaikwad <gaikwadabhishek1997@gmail.com>
Date: Thu, 6 Jun 2024 13:11:52 -0700
Subject: [PATCH] python: correct return type for raw stream in ObjectReader

Changed the return type of the raw() method in ObjectReader to return the correct file-like object instead of bytes. Improved docstrings and added type annotations for clarity.

Signed-off-by: Abhishek Gaikwad <gaikwadabhishek1997@gmail.com>
---
 python/aistore/sdk/object_reader.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/python/aistore/sdk/object_reader.py b/python/aistore/sdk/object_reader.py
index 497d26bde0..a197bb547b 100644
--- a/python/aistore/sdk/object_reader.py
+++ b/python/aistore/sdk/object_reader.py
@@ -1,8 +1,6 @@
 from typing import Iterator
-
 import requests
 from requests.structures import CaseInsensitiveDict
-
 from aistore.sdk.const import DEFAULT_CHUNK_SIZE
 from aistore.sdk.object_attributes import ObjectAttributes
 
@@ -10,7 +8,7 @@
 class ObjectReader:
     """
     Represents the data returned by the API when getting an object, including access to the content stream and object
-    attributes
+    attributes.
     """
 
     def __init__(
@@ -26,37 +24,42 @@ def __init__(
     @property
     def attributes(self) -> ObjectAttributes:
         """
-        Object metadata attributes
+        Object metadata attributes.
 
         Returns:
-            Object attributes parsed from the headers returned by AIS
+            ObjectAttributes: Parsed object attributes from the headers returned by AIS
         """
         return self._attributes
 
     def read_all(self) -> bytes:
         """
         Read all byte data from the object content stream.
-        This uses a bytes cast which makes it slightly slower and requires all object content to fit in memory at once
-        Returns:
-            Object content as bytes
 
+        This uses a bytes cast which makes it slightly slower and requires all object content to fit in memory at once.
+
+        Returns:
+            bytes: Object content as bytes.
         """
         obj_arr = bytearray()
         for chunk in self:
             obj_arr.extend(chunk)
         return bytes(obj_arr)
 
-    def raw(self) -> bytes:
+    def raw(self) -> requests.Response:
         """
-        Returns: Raw byte stream of object content
+        Returns the raw byte stream of object content.
+
+        Returns:
+            requests.Response: Raw byte stream of the object content
         """
         return self._stream.raw
 
     def __iter__(self) -> Iterator[bytes]:
         """
-        Creates a generator to read the stream content in chunks
+        Creates a generator to read the stream content in chunks.
+
         Returns:
-            An iterator with access to the next chunk of bytes
+            Iterator[bytes]: An iterator to access the next chunk of bytes
         """
         try:
             for chunk in self._stream.iter_content(chunk_size=self._chunk_size):