gdrive: add open

Fixes iterative#3408
Related iterative#2865
Fixes iterative#3897

Author: casperdcl

dvc/remote/gdrive.py

@@ -1,9 +1,11 @@
+import io
 import logging
 import os
 import posixpath
 import re
 import threading
 from collections import defaultdict
+from contextlib import contextmanager
 from urllib.parse import urlparse
 
 from funcy import cached_property, retry, wrap_prop, wrap_with
@@ -15,6 +17,7 @@
 from dvc.remote.base import BaseRemote
 from dvc.scheme import Schemes
 from dvc.utils import format_link, tmp_fname
+from dvc.utils.stream import IterStream
 
 logger = logging.getLogger(__name__)
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
@@ -393,6 +396,23 @@ def _gdrive_download_file(
         ) as pbar:
             gdrive_file.GetContentFile(to_file, callback=pbar.update_to)
 
+    @contextmanager
+    @_gdrive_retry
+    def open(self, path_info, mode="r", encoding=None):
+        assert mode in {"r", "rt", "rb"}
+
+        item_id = self._get_item_id(path_info)
+        param = {"id": item_id}
+        # it does not create a file on the remote
+        gdrive_file = self._drive.CreateFile(param)
+        fd = gdrive_file.GetContentIOBuffer()
+        stream = IterStream(iter(fd))
+
+        if mode != "rb":
+            stream = io.TextIOWrapper(stream, encoding=encoding)
+
+        yield stream
+
     @_gdrive_retry
     def _gdrive_delete_file(self, item_id):
         from pydrive2.files import ApiRequestError

dvc/utils/http.py

@@ -1,6 +1,8 @@
 import io
 from contextlib import contextmanager
 
+from dvc.utils.stream import IterStream
+
 
 @contextmanager
 def open_url(url, mode="r", encoding=None):
@@ -61,47 +63,3 @@ def gen(response):
     finally:
         # Ensure connection is closed
         it.close()
-
-
-class IterStream(io.RawIOBase):
-    """Wraps an iterator yielding bytes as a file object"""
-
-    def __init__(self, iterator):
-        self.iterator = iterator
-        self.leftover = None
-
-    def readable(self):
-        return True
-
-    # Python 3 requires only .readinto() method, it still uses other ones
-    # under some circumstances and falls back if those are absent. Since
-    # iterator already constructs byte strings for us, .readinto() is not the
-    # most optimal, so we provide .read1() too.
-
-    def readinto(self, b):
-        try:
-            n = len(b)  # We're supposed to return at most this much
-            chunk = self.leftover or next(self.iterator)
-            output, self.leftover = chunk[:n], chunk[n:]
-
-            n_out = len(output)
-            b[:n_out] = output
-            return n_out
-        except StopIteration:
-            return 0  # indicate EOF
-
-    readinto1 = readinto
-
-    def read1(self, n=-1):
-        try:
-            chunk = self.leftover or next(self.iterator)
-        except StopIteration:
-            return b""
-
-        # Return an arbitrary number or bytes
-        if n <= 0:
-            self.leftover = None
-            return chunk
-
-        output, self.leftover = chunk[:n], chunk[n:]
-        return output

dvc/utils/stream.py

@@ -0,0 +1,45 @@
+import io
+
+
+class IterStream(io.RawIOBase):
+    """Wraps an iterator yielding bytes as a file object"""
+
+    def __init__(self, iterator):
+        self.iterator = iterator
+        self.leftover = None
+
+    def readable(self):
+        return True
+
+    # Python 3 requires only .readinto() method, it still uses other ones
+    # under some circumstances and falls back if those are absent. Since
+    # iterator already constructs byte strings for us, .readinto() is not the
+    # most optimal, so we provide .read1() too.
+
+    def readinto(self, b):
+        try:
+            n = len(b)  # We're supposed to return at most this much
+            chunk = self.leftover or next(self.iterator)
+            output, self.leftover = chunk[:n], chunk[n:]
+
+            n_out = len(output)
+            b[:n_out] = output
+            return n_out
+        except StopIteration:
+            return 0  # indicate EOF
+
+    readinto1 = readinto
+
+    def read1(self, n=-1):
+        try:
+            chunk = self.leftover or next(self.iterator)
+        except StopIteration:
+            return b""
+
+        # Return an arbitrary number or bytes
+        if n <= 0:
+            self.leftover = None
+            return chunk
+
+        output, self.leftover = chunk[:n], chunk[n:]
+        return output

setup.py

@@ -87,7 +87,7 @@ def run(self):
 # Extra dependencies for remote integrations
 
 gs = ["google-cloud-storage==1.19.0"]
-gdrive = ["pydrive2>=1.4.13"]
+gdrive = ["pydrive2>=1.4.15"]
 s3 = ["boto3>=1.9.201"]
 azure = ["azure-storage-blob==2.1.0"]
 oss = ["oss2==2.6.1"]