Coverage for src/shephex/study/table/littletable_table.py: 100%

51 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-06-20 14:13 +0200

1from typing import Any, Dict, List 

2 

3from littletable import Table 

4 

5 

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

7 try: 

8 return getattr(obj, attr) 

9 except AttributeError: 

10 return None 

11 

12 

13class LittleTable: 

14 def __init__(self) -> None: 

15 self.table = Table() 

16 self.table.add_field('identifier', fn=None, default=None) 

17 self.table.add_field('status', fn=None, default=None) 

18 self.table.create_index('identifier') 

19 

20 self.column_names = ['identifier', 'status'] 

21 

22 def add_column(self, name: str) -> None: 

23 """ 

24 Add a new column to the table. The column will be filled with None values. 

25 """ 

26 self.table.add_field(name, fn=None, default=None) 

27 self.column_names.append(name) 

28 

29 def add_row(self, row_data: Dict, add_columns: bool = False) -> None: 

30 for key in row_data.keys(): 

31 if key not in self.column_names: 

32 self.add_column(key) 

33 

34 self.table.insert(row_data) 

35 

36 def contains_row(self, row_data: Dict) -> bool: 

37 """ 

38 Check if the table contains a row with the same data. 

39 """ 

40 # Check for the identifier 

41 contains_id = len(self.table.where(identifier=row_data['identifier'])) > 0 

42 if contains_id: 

43 return True 

44 

45 # Check if any row contains the same data 

46 unchecked_keys = [ 

47 'identifier', 

48 'status', 

49 'time_stamp', 

50 'procedure_path', 

51 'options_path', 

52 'procedure' 

53 ] 

54 checked_row = { 

55 key: row_data[key] for key in row_data if key not in unchecked_keys 

56 } 

57 contains_data = len(self.table.where(**checked_row)) > 0 

58 

59 if contains_data: 

60 return True 

61 

62 return False 

63 

64 def get_row_match(self, identifier: str) -> Dict: 

65 # First find the corresponding row with the identifier 

66 match_list = self.table.where(identifier=identifier) 

67 if len(match_list) > 1 or len(match_list) == 0: 

68 raise ValueError( 

69 f'Found {len(match_list)} rows with the identifier {identifier}.' 

70 ) 

71 

72 match = match_list[0] 

73 return match 

74 

75 def update_row(self, row_data: Dict) -> None: 

76 """ 

77 Update the row in the table. 

78 """ 

79 match = self.get_row_match(row_data['identifier']) 

80 self.table.remove(match) 

81 self.add_row(row_data) 

82 

83 def update_row_partially(self, row_data: Dict) -> None: 

84 match = self.get_row_match(row_data['identifier']) 

85 for key in self.column_names: 

86 if key not in row_data.keys(): 

87 row_data[key] = safe_getattr(match, key) 

88 

89 self.table.remove(match) 

90 self.add_row(row_data) 

91 

92 def where(self, *args, **kwargs) -> List[Dict]: 

93 """ 

94 Return identifiers of rows that match the query. 

95 """ 

96 return [row.identifier for row in self.table.where(*args, **kwargs)]