add multi_apps; update chrome utilities

desktop_env/evaluators/getters/__init__.py

@@ -1,7 +1,6 @@
 from .chrome import get_default_search_engine, get_cookie_data, get_bookmarks, get_open_tabs_info, get_pdf_from_url, \
     get_shortcuts_on_desktop, get_history, get_enabled_experiments, get_chrome_language, get_chrome_font_size, \
-    get_profile_name, \
-    get_number_of_search_results
+    get_profile_name, get_number_of_search_results, get_googledrive_file
 from .file import get_cloud_file, get_vm_file, get_cache_file
 from .general import get_vm_command_line
 from .impress import get_audio_in_slide
@@ -9,4 +8,4 @@ from .info import get_vm_screen_size, get_vm_window_size, get_vm_wallpaper
 from .misc import get_rule, get_accessibility_tree
 from .replay import get_replay
 from .vlc import get_vlc_playing_info, get_vlc_config
-from .vscode import get_vscode_config
+from .vscode import get_vscode_config

desktop_env/evaluators/getters/chrome.py

@@ -2,8 +2,9 @@ import json
 import logging
 import os
 import sqlite3
-from typing import Dict
-
+from typing import Dict, Any
+from pydrive.auth import GoogleAuth
+from pydrive.drive import GoogleDrive, GoogleDriveFileList, GoogleDriveFile
 from playwright.sync_api import sync_playwright
 
 logger = logging.getLogger("desktopenv.getters.chrome")
@@ -461,3 +462,49 @@ def get_number_of_search_results(env, config: Dict[str, str]):
         browser.close()
 
     return actual_count
+
+
+def get_googledrive_file(env, config: Dict[str, Any]) -> str:
+    """ Get the desired file from Google Drive based on config, return the downloaded local filepath.
+    """
+    settings_file = config.get('settings_file', 'evaluation_examples/settings/googledrive/settings.json')
+    auth = GoogleAuth(settings_file=settings_file)
+    drive = GoogleDrive(auth)
+
+    q = config['query']
+    filelist: GoogleDriveFileList = drive.ListFile({'q': q}).GetList()
+    if len(filelist) == 0: # target file not found
+        return None
+    
+    file: GoogleDriveFile = filelist[0] # HACK: if multiple candidates, just download the first one
+    _path = os.path.join(env.cache_dir, config['dest'])
+
+    try:
+        file.GetContentFile(_path, mimetype=file.metadata['mimeType'])
+    except:
+        logger.info('[ERROR]: Failed to download the file from Google Drive')
+        return None
+    return _path
+
+
+def get_googledrive_file(env, config: Dict[str, Any]) -> str:
+    """ Get the desired file from Google Drive based on config, return the downloaded local filepath.
+    """
+    settings_file = config.get('settings_file', 'evaluation_examples/settings/googledrive/settings.json')
+    auth = GoogleAuth(settings_file=settings_file)
+    drive = GoogleDrive(auth)
+
+    q = config['query']
+    filelist: GoogleDriveFileList = drive.ListFile({'q': q}).GetList()
+    if len(filelist) == 0: # target file not found
+        return None
+    
+    file: GoogleDriveFile = filelist[0] # HACK: if multiple candidates, just download the first one
+    _path = os.path.join(env.cache_dir, config['dest'])
+
+    try:
+        file.GetContentFile(_path, mimetype=file.metadata['mimeType'])
+    except:
+        logger.info('[ERROR]: Failed to download the file from Google Drive')
+        return None
+    return _path