diff --git a/youtube_cli/main.py b/youtube_cli/main.py index da872f9..3a1c1c6 100644 --- a/youtube_cli/main.py +++ b/youtube_cli/main.py @@ -661,17 +661,35 @@ class YouTubeCLI: while True: try: - console.print("\n[blue]Select a category (enter number):[/blue]") + console.print("\n[blue]Select a category (enter number) or type a custom folder name:[/blue]") choice = input().strip() - choice_num = int(choice) - if 1 <= choice_num <= len(categories): - selected_category = categories[choice_num - 1] - console.print(f"[green]Selected category: {Path(selected_category).name}[/green]") - return selected_category + + # Check if input is a number + if choice.isdigit(): + choice_num = int(choice) + if 1 <= choice_num <= len(categories): + selected_category = categories[choice_num - 1] + console.print(f"[green]Selected category: {Path(selected_category).name}[/green]") + return selected_category + else: + console.print("[red]Invalid selection. Please try again.[/red]") else: - console.print("[red]Invalid selection. Please try again.[/red]") - except ValueError: - console.print("[red]Please enter a valid number.[/red]") + # Validate custom folder name + if not choice: + console.print("[red]Folder name cannot be empty. Please try again.[/red]") + continue + + # Check for invalid characters (only allow a-z, 0-9, and hyphens/underscores) + import re + if not re.match(r'^[a-zA-Z0-9_-]+$', choice): + console.print("[red]Invalid characters. Only a-z, 0-9, hyphens, and underscores are allowed.[/red]") + continue + + # If valid, use the custom folder name + console.print(f"[green]Using custom folder: {choice}[/green]") + # Return the custom folder name (will be appended to base path) + return choice + except KeyboardInterrupt: console.print("\n[yellow]Operation cancelled.[/yellow]") return None