Coverage for src/shephex/study/renderer.py: 93%

54 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1from typing import Any 

2 

3from rich.console import Console 

4from rich.table import Table 

5 

6from shephex.study.study import Study 

7 

8 

9def safe_getattr(obj: Any, attr: str) -> Any: 

10 try: 

11 return getattr(obj, attr) 

12 except AttributeError: 

13 return None 

14 

15class StudyRenderer: 

16 

17 def __init__(self, excluded: list[str] = None) -> None: 

18 self.excluded = ['procedure', 'options_path', 'time_stamp', 'procedure_type', 'procedure_path'] 

19 self.capitalize_fields = ['identifier', 'status'] 

20 if excluded is not None: 

21 self.excluded += excluded 

22 self.conditions = {} 

23 

24 def add_condition(self, **kwargs) -> None: 

25 self.conditions.update(kwargs) 

26 

27 def initialize_table(self, study: Study, **kwargs) -> None: 

28 rich_table = Table(**kwargs) 

29 for field in study.table.column_names: 

30 if field in self.excluded: 

31 continue 

32 if field in self.capitalize_fields: 

33 field = field.capitalize() 

34 

35 rich_table.add_column(field, justify='center') 

36 return rich_table 

37 

38 def get_rows(self, study: Study) -> list: 

39 # Janky code to sort by time stamp. 

40 rows = [row for row in study.table.table] 

41 time_stamps = [row.time_stamp for row in rows] 

42 sorted_indices = sorted(range(len(time_stamps)), key=lambda k: time_stamps[k]) 

43 rows = [rows[i] for i in sorted_indices] 

44 return rows 

45 

46 def get_table(self, study: Study, **kwargs) -> Table: 

47 style_dict = { 

48 'pending': 'yellow', 

49 'running': 'blue', 

50 'completed': 'green', 

51 'failed': 'red', 

52 } 

53 

54 rich_table = self.initialize_table(study, **kwargs) 

55 rows = self.get_rows(study) 

56 allowed_identifiers = study.table.where(**self.conditions) 

57 

58 # Add rows to the rich.Table 

59 for row in rows: 

60 if row.identifier not in allowed_identifiers: 

61 continue 

62 

63 row_data = [] 

64 for name in study.table.column_names: 

65 if name in self.excluded: 

66 continue 

67 if name == 'status': 

68 row_data.append(row.status.capitalize()) 

69 else: 

70 row_data.append(str(safe_getattr(row, name))) 

71 rich_table.add_row( 

72 *row_data, style=style_dict.get(row.status.lower(), 'white') 

73 ) 

74 

75 return rich_table 

76 

77 def render_study(self, study: Study, **kwargs) -> None: 

78 rich_table = self.get_table(study, **kwargs) 

79 

80 # Decide whether to return the table or print # noqa 

81 console = Console() 

82 console.print(rich_table)