Coverage for /builds/kinetik161/ase/ase/db/web.py: 72.90%

1"""Helper functions for Flask WSGI-app.""" 

2from typing import Dict, List, Optional, Tuple 

3 

4from ase.db.core import Database 

5from ase.db.table import Table, all_columns 

6 

7 

8class Session: 

9 """Seesion object. 

10 

11 Stores stuff that the jinja2 templetes (like templates/table.html) 

12 need to show. Example from table.html:: 

13 

14 Displaying rows {{ s.row1 }}-{{ s.row2 }} out of {{ s.nrows }} 

15 

16 where *s* is the session object. 

17 """ 

18 next_id = 1 

19 sessions: Dict[int, 'Session'] = {} 

20 

21 def __init__(self, project_name: str): 

22 self.id = Session.next_id 

23 Session.next_id += 1 

24 

25 Session.sessions[self.id] = self 

26 if len(Session.sessions) > 2000: 

27 # Forget old sessions: 

28 for id in sorted(Session.sessions)[:400]: 

29 del Session.sessions[id] 

30 

31 self.columns: Optional[List[str]] = None 

32 self.nrows: Optional[int] = None 

33 self.nrows_total: Optional[int] = None 

34 self.page = 0 

35 self.limit = 25 

36 self.sort = '' 

37 self.query = '' 

38 self.project_name = project_name 

39 

40 def __str__(self) -> str: 

41 return str(self.__dict__) 

42 

43 @staticmethod 

44 def get(id: int) -> 'Session': 

45 return Session.sessions[id] 

46 

47 def update(self, 

48 what: str, 

49 x: str, 

50 args: Dict[str, str], 

51 project) -> None: 

52 

53 if self.columns is None: 

54 self.columns = list(project.default_columns) 

55 

56 if what == 'query': 

57 self.query = project.handle_query(args) 

58 self.nrows = None 

59 self.page = 0 

60 

61 elif what == 'sort': 

62 if x == self.sort: 

63 self.sort = '-' + x 

64 elif '-' + x == self.sort: 

65 self.sort = 'id' 

66 else: 

67 self.sort = x 

68 self.page = 0 

69 

70 elif what == 'limit': 

71 self.limit = int(x) 

72 self.page = 0 

73 

74 elif what == 'page': 

75 self.page = int(x) 

76 

77 elif what == 'toggle': 

78 column = x 

79 if column == 'reset': 

80 self.columns = list(project.default_columns) 

81 else: 

82 if column in self.columns: 

83 self.columns.remove(column) 

84 if column == self.sort.lstrip('-'): 

85 self.sort = 'id' 

86 self.page = 0 

87 else: 

88 self.columns.append(column) 

89 

90 @property 

91 def row1(self) -> int: 

92 return self.page * self.limit + 1 

93 

94 @property 

95 def row2(self) -> int: 

96 assert self.nrows is not None 

97 return min((self.page + 1) * self.limit, self.nrows) 

98 

99 def paginate(self) -> List[Tuple[int, str]]: 

100 """Helper function for pagination stuff.""" 

101 assert self.nrows is not None 

102 npages = (self.nrows + self.limit - 1) // self.limit 

103 p1 = min(5, npages) 

104 p2 = max(self.page - 4, p1) 

105 p3 = min(self.page + 5, npages) 

106 p4 = max(npages - 4, p3) 

107 pgs = list(range(p1)) 

108 if p1 < p2: 

109 pgs.append(-1) 

110 pgs += list(range(p2, p3)) 

111 if p3 < p4: 

112 pgs.append(-1) 

113 pgs += list(range(p4, npages)) 

114 pages = [(self.page - 1, 'previous')] 

115 for p in pgs: 

116 if p == -1: 

117 pages.append((-1, '...')) 

118 elif p == self.page: 

119 pages.append((-1, str(p + 1))) 

120 else: 

121 pages.append((p, str(p + 1))) 

122 nxt = min(self.page + 1, npages - 1) 

123 if nxt == self.page: 

124 nxt = -1 

125 pages.append((nxt, 'next')) 

126 return pages 

127 

128 def create_table(self, 

129 db: Database, 

130 uid_key: str, 

131 keys: List[str]) -> Table: 

132 query = self.query 

133 

134 if self.nrows_total is None: 

135 self.nrows_total = db.count() 

136 

137 if self.nrows is None: 

138 try: 

139 self.nrows = db.count(query) 

140 except (ValueError, KeyError) as e: 

141 error = ', '.join(['Bad query'] + list(e.args)) 

142 from flask import flash 

143 flash(error) 

144 query = 'id=0' # this will return no rows 

145 self.nrows = 0 

146 

147 table = Table(db, uid_key) 

148 table.select(query, self.columns, self.sort, 

149 self.limit, offset=self.page * self.limit, 

150 show_empty_columns=True) 

151 table.format() 

152 assert self.columns is not None 

153 table.addcolumns = sorted(column for column in 

154 [*all_columns, *keys] 

155 if column not in self.columns) 

156 return table