fix(GoogleDriveFile): Add remote file helper for GDrive

trame/__init__.py

@@ -1,7 +1,7 @@
 from trame.internal import (
   change, Controller, flush_state, get_cli_parser, get_state, get_version,
   is_dirty, is_dirty_all, port, setup_dev, start, State, stop, trigger,
-  update_state, RemoteFile, Singleton
+  update_state, RemoteFile, Singleton, GoogleDriveFile
 )
 from trame.layouts import update_layout
 
@@ -86,6 +86,7 @@
 
     # Utils/helpers
     "RemoteFile",
+    "GoogleDriveFile",
 
     # These are not exposed in the docs
     "__version__",

trame/internal/__init__.py

@@ -8,5 +8,5 @@
 from .utils import (
     AppServerThread, base_directory, ClientWindowProcess, compose_callbacks,
     get_cli_parser, get_version, is_dunder, log_js_error, print_server_info,
-    validate_key_names, RemoteFile, Singleton
+    validate_key_names, RemoteFile, Singleton, GoogleDriveFile
 )

trame/internal/utils/__init__.py

@@ -5,7 +5,7 @@
 from .logging import log_js_error, print_server_info, validate_key_names
 from .string import is_dunder
 from .version import get_version
-from .remote_data import RemoteFile
+from .remote_data import RemoteFile, GoogleDriveFile
 from .singleton import Singleton
 
 __all__ = [
@@ -20,5 +20,6 @@
     "print_server_info",
     "validate_key_names",
     "RemoteFile",
+    "GoogleDriveFile",
     "Singleton",
 ]

trame/internal/utils/remote_data.py

@@ -3,9 +3,41 @@
 from urllib.request import urlretrieve
 
 
-class RemoteFile:
-    def __init__(self, local_path=None, remote_url=None, local_base=None):
-        self._url = remote_url
+def download_file_from_google_drive(id, destination):
+    import requests
+
+    URL = "https://docs.google.com/uc?export=download"
+
+    session = requests.Session()
+    response = session.get(URL, params={"id": id}, stream=True)
+    token = get_confirm_token(response)
+
+    if token:
+        params = {"id": id, "confirm": token}
+        response = session.get(URL, params=params, stream=True)
+
+    save_response_content(response, destination)
+
+
+def get_confirm_token(response):
+    for key, value in response.cookies.items():
+        if key.startswith("download_warning"):
+            return value
+
+    return None
+
+
+def save_response_content(response, destination):
+    CHUNK_SIZE = 32768
+
+    with open(destination, "wb") as f:
+        for chunk in response.iter_content(CHUNK_SIZE):
+            if chunk:  # filter out keep-alive new chunks
+                f.write(chunk)
+
+
+class AbstractRemoteFile:
+    def __init__(self, local_path=None, local_base=None):
         # setup base path
         self._base = os.getcwd()
         if local_base is not None:
@@ -32,15 +64,37 @@ def local(self):
         return os.path.exists(self._file_path)
 
     def fetch(self):
-        try:
-            print(f"Downloading:\n - {self._url}\n - to {self._file_path}")
-            urlretrieve(self._url, self._file_path)
-        except HTTPError as e:
-            print(RuntimeError(f"Failed to download {self._url}. {e.reason}"))
+        pass
 
     @property
     def path(self):
         if not self.local:
             self.fetch()
 
         return self._file_path
+
+
+class GoogleDriveFile(AbstractRemoteFile):
+    def __init__(self, local_path=None, google_id=None, local_base=None):
+        super().__init__(local_path, local_base)
+        self._gid = google_id
+
+    def fetch(self):
+        try:
+            print(f"Downloading:\n - {self._gid}\n - to {self._file_path}")
+            download_file_from_google_drive(self._gid, self._file_path)
+        except HTTPError as e:
+            print(RuntimeError(f"Failed to download {self._gid}. {e.reason}"))
+
+
+class RemoteFile(AbstractRemoteFile):
+    def __init__(self, local_path=None, remote_url=None, local_base=None):
+        super().__init__(local_path, local_base)
+        self._url = remote_url
+
+    def fetch(self):
+        try:
+            print(f"Downloading:\n - {self._url}\n - to {self._file_path}")
+            urlretrieve(self._url, self._file_path)
+        except HTTPError as e:
+            print(RuntimeError(f"Failed to download {self._url}. {e.reason}"))