add len(MediaIoReadable) (#40)

* add len(MediaIoReadable) * test buffer len
iterative · May 31, 2020 · 9275547 · 9275547
1 parent d1ef328
commit 9275547
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/pydrive2/files.py b/pydrive2/files.py
@@ -137,6 +137,7 @@ def __init__(
         self.downloader = MediaIoBaseDownload(
             self._fd, request, chunksize=chunksize
         )
+        self.size = None
         self._pre_buffer = False
         if pre_buffer:
             self.read()
@@ -157,7 +158,8 @@ def read(self):
         if self.done:
             return None
         try:
-            _, self.done = self.downloader.next_chunk()
+            status, self.done = self.downloader.next_chunk()
+            self.size = status.total_size
         except errors.HttpError as error:
             raise ApiRequestError(error)
         return self._fd.read()
@@ -172,6 +174,9 @@ def __iter__(self):
                 break
             yield chunk
 
+    def __len__(self):
+        return self.size
+
 
 class GoogleDriveFile(ApiAttributeMixin, ApiResource):
     """Google Drive File instance.

diff --git a/pydrive2/test/test_file.py b/pydrive2/test/test_file.py
@@ -291,9 +291,11 @@ def test_11_Files_Get_Content_Buffer(self):
 
         buffer1 = pydrive_retry(file1.GetContentIOBuffer)
         self.assertEqual(file1.metadata["title"], filename)
+        self.assertEqual(len(buffer1), len(content))
         self.assertEqual(b"".join(iter(buffer1)).decode("ascii"), content)
 
         buffer2 = pydrive_retry(file1.GetContentIOBuffer, encoding="ascii")
+        self.assertEqual(len(buffer2), len(content))
         self.assertEqual("".join(iter(buffer2)), content)
 
         self.DeleteUploadedFiles(drive, [file1["id"]])