Coverage for src/utils/ui.py: 35%

1# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL 

2# Unauthorized copying of this file and/or project, via any medium is strictly prohibited. 

3# Direct all questions to legal@cloudzero.com 

4 

5""" 

6UI and display utilities for CloudZero Data Tool. 

7Handles user interaction, table display, and console output. 

8""" 

9 

10import polars as pl 

11 

12try: 

13 from .table_helpers import create_basic_table, TableColumn 

14 from .console_helpers import console 

15except ImportError: 

16 # Fallback for direct execution or test imports 

17 import sys 

18 from pathlib import Path 

19 sys.path.append(str(Path(__file__).parent)) 

20 from table_helpers import create_basic_table, TableColumn 

21 from console_helpers import console 

22 

23 

24def show_column_info(df: pl.DataFrame) -> None: 

25 """Display information about all columns in the DataFrame.""" 

26 console.print("\n📊 [bold blue]Column Information[/bold blue]") 

27 

28 columns = [ 

29 TableColumn("Column", "cyan"), 

30 TableColumn("Type", "magenta"), 

31 TableColumn("Non-null Count", "green", justify="right"), 

32 TableColumn("Sample Values", "yellow") 

33 ] 

34 table = create_basic_table("DataFrame Schema", columns) 

35 

36 for col in df.columns: 

37 dtype = str(df[col].dtype) 

38 non_null_count = df[col].drop_nulls().len() 

39 

40 # Get sample values (first few non-null values) 

41 sample_values = df[col].drop_nulls().head(3).to_list() 

42 sample_str = ", ".join( 

43 [str(v)[:20] + "..." if len(str(v)) > 20 else str(v) for v in sample_values] 

44 ) 

45 

46 table.add_row(col, dtype, f"{non_null_count:,}", sample_str) 

47 

48 console.print(table) 

49 

50 

51def get_column_choice(df: pl.DataFrame) -> str: 

52 """Get user's choice of column to analyze.""" 

53 columns = df.columns 

54 

55 console.print("\n🔍 [bold blue]Available Columns:[/bold blue]") 

56 for i, col in enumerate(columns, 1): 

57 console.print(f" {i:2d}. {col}") 

58 

59 while True: 

60 try: 

61 choice = console.input( 

62 f"\n[bold green]Enter column number (1-{len(columns)}) or column name: [/bold green]" 

63 ).strip() 

64 

65 # Try to parse as number first 

66 try: 

67 col_num = int(choice) 

68 if 1 <= col_num <= len(columns): 

69 return columns[col_num - 1] 

70 else: 

71 console.print( 

72 f"[red]Please enter a number between 1 and {len(columns)}[/red]" 

73 ) 

74 continue 

75 except ValueError: 

76 # Try to match by name 

77 if choice in columns: 

78 return choice 

79 else: 

80 # Try case-insensitive match 

81 matches = [col for col in columns if col.lower() == choice.lower()] 

82 if matches: 

83 return matches[0] 

84 else: 

85 console.print( 

86 f"[red]Column '{choice}' not found. Please try again.[/red]" 

87 ) 

88 continue 

89 

90 except KeyboardInterrupt: 

91 console.print("\n[yellow]Goodbye![/yellow]") 

92 exit(0)