| 201 | |
| 202 | |
| 203 | class ColabEnvironment(IPythonZMQEnvironment): |
| 204 | name = "Google Colab Environment" |
| 205 | |
| 206 | def get_notebook_json(self) -> dict: |
| 207 | """Get the JSON for the current Colab notebook""" |
| 208 | import ipynbname |
| 209 | from google.colab import auth |
| 210 | from googleapiclient.discovery import build |
| 211 | from googleapiclient.http import MediaIoBaseDownload |
| 212 | |
| 213 | # Get the notebook's Google Drive file_id |
| 214 | file_id = ipynbname.name().replace("fileId=", "") |
| 215 | |
| 216 | try: |
| 217 | auth.authenticate_user() |
| 218 | except Exception as e: |
| 219 | raise NotebookException( |
| 220 | "Google Drive authentication failed. Please allow this notebook to access your Google Drive." |
| 221 | ) from e |
| 222 | |
| 223 | drive_service = build("drive", "v3") |
| 224 | |
| 225 | request = drive_service.files().get_media(fileId=file_id) |
| 226 | downloaded = io.BytesIO() |
| 227 | downloader = MediaIoBaseDownload(downloaded, request) |
| 228 | done = False |
| 229 | while done is False: |
| 230 | # _ is a placeholder for a progress object that we ignore. |
| 231 | # (Our file is small, so we skip reporting progress.) |
| 232 | _, done = downloader.next_chunk() |
| 233 | |
| 234 | downloaded.seek(0) |
| 235 | notebook_json = json.loads(downloaded.read().decode("utf-8")) |
| 236 | |
| 237 | return notebook_json |
| 238 | |
| 239 | |
| 240 | class CodespacesVSCodeJupyterEnvironment(VSCodeJupyterEnvironment): |