11import typer
2+ import os
3+ from pathlib import Path
24from rich .console import Console
35from rich .table import Table
46from rich import print
5- from rich .progress import Progress
7+ from rich .progress import Progress , SpinnerColumn , TextColumn , BarColumn , DownloadColumn , TransferSpeedColumn
68from rich .panel import Panel
79from rich .text import Text
810from urllib .parse import urlparse
@@ -161,6 +163,75 @@ def list_artifacts(job_id: str, output_format: str = "pretty") -> None:
161163 console .print (f"[red]Error:[/red] Failed to fetch artifacts. Status code: { response .status_code } )
162164
163165
166+ def download_artifacts (job_id : str , output_dir : str = None ) -> None :
167+ """Download all artifacts for a job as a zip file."""
168+ if output_dir is None :
169+ output_dir = os .getcwd ()
170+ else :
171+ output_dir = os .path .abspath (output_dir )
172+ Path (output_dir ).mkdir (parents = True , exist_ok = True )
173+
174+ # Determine output filename
175+ filename = f"artifacts_{ job_id }
176+ output_path = os .path .join (output_dir , filename )
177+
178+ # Check if file already exists
179+ if os .path .exists (output_path ):
180+ console .print (f"[yellow]Warning:[/yellow] File { output_path } )
181+
182+ try :
183+ with console .status (f"[bold green]Downloading artifacts for job { job_id } , spinner = "dots" ):
184+ response = api .get (f"/jobs/{ job_id } , timeout = 300.0 )
185+
186+ if response .status_code == 200 :
187+ # Get filename from Content-Disposition header if available
188+ content_disposition = response .headers .get ("Content-Disposition" , "" )
189+ if "filename=" in content_disposition :
190+ # Extract filename from Content-Disposition header
191+ filename_part = content_disposition .split ("filename=" )[1 ].strip ('"' )
192+ if filename_part :
193+ filename = filename_part
194+ output_path = os .path .join (output_dir , filename )
195+
196+ # Write the file with progress tracking
197+ content_length = response .headers .get ("Content-Length" )
198+ total_size = int (content_length ) if content_length else None
199+
200+ with Progress (
201+ SpinnerColumn (),
202+ TextColumn ("[progress.description]{task.description}" ),
203+ BarColumn (),
204+ DownloadColumn (),
205+ TransferSpeedColumn (),
206+ ) as progress :
207+ task = progress .add_task (f"[cyan]Downloading { filename } , total = total_size )
208+
209+ with open (output_path , "wb" ) as f :
210+ if total_size :
211+ # Show progress if we know the total size
212+ downloaded = 0
213+ for chunk in response .iter_bytes (chunk_size = 8192 ):
214+ if chunk :
215+ f .write (chunk )
216+ downloaded += len (chunk )
217+ progress .update (task , completed = downloaded )
218+ else :
219+ # Just write without progress if we don't know the size
220+ f .write (response .content )
221+ progress .update (task , completed = 1 , total = 1 )
222+
223+ console .print (f"[green]✓[/green] Successfully downloaded artifacts to: { output_path } )
224+ elif response .status_code == 404 :
225+ console .print (f"[red]Error:[/red] No artifacts found for job { job_id } )
226+ else :
227+ console .print (f"[red]Error:[/red] Failed to download artifacts. Status code: { response .status_code } )
228+ if response .text :
229+ console .print (f"[red]Response:[/red] { response .text [:200 ]} )
230+
231+ except Exception as e :
232+ console .print (f"[red]Error:[/red] Failed to download artifacts: { e } )
233+
234+
164235@app .command ("artifacts" )
165236def command_job_artifacts (
166237 job_id : str = typer .Argument (..., help = "Job ID to list artifacts for" ),
@@ -170,6 +241,18 @@ def command_job_artifacts(
170241 list_artifacts (job_id )
171242
172243
244+ @app .command ("download" )
245+ def command_job_download (
246+ job_id : str = typer .Argument (..., help = "Job ID to download artifacts for" ),
247+ output_dir : str = typer .Option (
248+ None , "--output" , "-o" , help = "Output directory for the zip file (default: current directory)"
249+ ),
250+ ):
251+ """Download all artifacts for a job as a zip file."""
252+ check_configs ()
253+ download_artifacts (job_id , output_dir )
254+
255+
173256@app .command ("list" )
174257def command_job_list ():
175258 """List all jobs."""
0 commit comments