| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | import click |
| 16 | import hashlib |
| 17 | import os |
| 18 | import tempfile |
| 19 | import typer |
| 20 | from typing import Optional |
| 21 | from typing_extensions import Annotated |
| 22 | |
| 23 | from colab_cli.contents import ContentsClient |
| 24 | |
| 25 | |
| 26 | def ls( |
| 27 | session: Annotated[ |
| 28 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 29 | ] = None, |
| 30 | path: Annotated[str, typer.Argument(help="Remote path to list")] = "content", |
| 31 | ): |
| 32 | """List files in a session""" |
| 33 | from colab_cli.common import state |
| 34 | |
| 35 | name = state.resolve_session(session) |
| 36 | s = state.store.get(name) |
| 37 | if not s: |
| 38 | typer.echo(f"[colab] Session '{name}' not found.") |
| 39 | raise typer.Exit(1) |
| 40 | contents = ContentsClient(s) |
| 41 | try: |
| 42 | data = contents.list_dir(path) |
| 43 | state.history.log_event(name, "file_operation", {"op": "ls", "path": path}) |
| 44 | if data.get("type") == "directory": |
| 45 | items = data.get("content", []) |
| 46 | for item in sorted( |
| 47 | items, key=lambda x: (x.get("type") != "directory", x.get("name")) |
| 48 | ): |
| 49 | suffix = "/" if item.get("type") == "directory" else "" |
| 50 | typer.echo(f"{item.get('name')}{suffix}") |
| 51 | else: |
| 52 | typer.echo(data.get("name")) |
| 53 | except Exception as e: |
| 54 | typer.echo(f"[colab] Error: {e}") |
| 55 | raise typer.Exit(1) |
| 56 | |
| 57 | |
| 58 | def rm( |
| 59 | session: Annotated[ |
| 60 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 61 | ] = None, |
| 62 | path: Annotated[str, typer.Argument(help="Remote path to remove")] = ..., |
| 63 | ): |
| 64 | """Remove a remote file""" |
| 65 | from colab_cli.common import state |
| 66 | |
| 67 | name = state.resolve_session(session) |
| 68 | s = state.store.get(name) |
| 69 | if not s: |
| 70 | typer.echo(f"[colab] Session '{name}' not found.") |
| 71 | raise typer.Exit(1) |
| 72 | contents = ContentsClient(s) |
| 73 | try: |
| 74 | contents.rm(path) |
| 75 | state.history.log_event(name, "file_operation", {"op": "rm", "path": path}) |
| 76 | typer.echo(f"[colab] Deleted {path}") |
| 77 | except Exception as e: |
| 78 | typer.echo(f"[colab] Error: {e}") |
| 79 | raise typer.Exit(1) |
| 80 | |
| 81 | |
| 82 | def upload( |
| 83 | session: Annotated[ |
| 84 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 85 | ] = None, |
| 86 | local_path: Annotated[str, typer.Argument(help="Local file to upload")] = ..., |
| 87 | remote_path: Annotated[str, typer.Argument(help="Remote path to upload to")] = ..., |
| 88 | ): |
| 89 | """Upload a file to a session""" |
| 90 | from colab_cli.common import state |
| 91 | |
| 92 | name = state.resolve_session(session) |
| 93 | s = state.store.get(name) |
| 94 | if not s: |
| 95 | typer.echo(f"[colab] Session '{name}' not found.") |
| 96 | raise typer.Exit(1) |
| 97 | if not os.path.isfile(local_path): |
| 98 | typer.echo(f"[colab] Local file '{local_path}' not found.") |
| 99 | raise typer.Exit(1) |
| 100 | contents = ContentsClient(s) |
| 101 | try: |
| 102 | contents.upload(local_path, remote_path) |
| 103 | state.history.log_event( |
| 104 | name, |
| 105 | "file_operation", |
| 106 | {"op": "upload", "local": local_path, "remote": remote_path}, |
| 107 | ) |
| 108 | typer.echo(f"[colab] Uploaded '{local_path}' to '{remote_path}'") |
| 109 | except Exception as e: |
| 110 | typer.echo(f"[colab] Upload failed: {e}") |
| 111 | raise typer.Exit(1) |
| 112 | |
| 113 | |
| 114 | def download( |
| 115 | session: Annotated[ |
| 116 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 117 | ] = None, |
| 118 | remote_path: Annotated[ |
| 119 | str, typer.Argument(help="Remote path to download from") |
| 120 | ] = ..., |
| 121 | local_path: Annotated[ |
| 122 | str, typer.Argument(help="Local path to save the file") |
| 123 | ] = ..., |
| 124 | ): |
| 125 | """Download a file from a session""" |
| 126 | from colab_cli.common import state |
| 127 | |
| 128 | name = state.resolve_session(session) |
| 129 | s = state.store.get(name) |
| 130 | if not s: |
| 131 | typer.echo(f"[colab] Session '{name}' not found.") |
| 132 | raise typer.Exit(1) |
| 133 | contents = ContentsClient(s) |
| 134 | try: |
| 135 | contents.download(remote_path, local_path) |
| 136 | state.history.log_event( |
| 137 | name, |
| 138 | "file_operation", |
| 139 | {"op": "download", "remote": remote_path, "local": local_path}, |
| 140 | ) |
| 141 | typer.echo(f"[colab] Downloaded '{remote_path}' to '{local_path}'") |
| 142 | except Exception as e: |
| 143 | typer.echo(f"[colab] Download failed: {e}") |
| 144 | raise typer.Exit(1) |
| 145 | |
| 146 | |
| 147 | def edit( |
| 148 | session: Annotated[ |
| 149 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 150 | ] = None, |
| 151 | remote_path: Annotated[str, typer.Argument(help="Remote path to edit")] = ..., |
| 152 | ): |
| 153 | """Edit a file on a running Colab session""" |
| 154 | from colab_cli.common import state |
| 155 | |
| 156 | name = state.resolve_session(session) |
| 157 | s = state.store.get(name) |
| 158 | if not s: |
| 159 | typer.echo(f"[colab] Session '{name}' not found.") |
| 160 | raise typer.Exit(1) |
| 161 | |
| 162 | contents = ContentsClient(s) |
| 163 | |
| 164 | def get_file_hash(path): |
| 165 | if not os.path.exists(path): |
| 166 | return None |
| 167 | with open(path, "rb") as f: |
| 168 | return hashlib.file_digest(f, "sha256").hexdigest() |
| 169 | |
| 170 | _, ext = os.path.splitext(remote_path) |
| 171 | |
| 172 | with tempfile.NamedTemporaryFile(suffix=ext) as tf: |
| 173 | local_path = tf.name |
| 174 | |
| 175 | try: |
| 176 | contents.download(remote_path, local_path) |
| 177 | except Exception: |
| 178 | # If download fails, assume file doesn't exist and start empty |
| 179 | pass |
| 180 | |
| 181 | hash_before = get_file_hash(local_path) |
| 182 | |
| 183 | click.edit(filename=local_path) |
| 184 | |
| 185 | hash_after = get_file_hash(local_path) |
| 186 | |
| 187 | if hash_after != hash_before: |
| 188 | contents.upload(local_path, remote_path) |
| 189 | state.history.log_event( |
| 190 | name, |
| 191 | "file_operation", |
| 192 | {"op": "edit", "remote": remote_path}, |
| 193 | ) |
| 194 | typer.echo(f"[colab] Edited and uploaded '{remote_path}'") |
| 195 | else: |
| 196 | typer.echo(f"[colab] No changes made to '{remote_path}'") |
| 197 | |
| 198 | |
| 199 | def register(app: typer.Typer): |
| 200 | app.command()(ls) |
| 201 | app.command()(rm) |
| 202 | app.command()(upload) |
| 203 | app.command()(download) |
| 204 | app.command()(edit) |