Coverage for collectra/gui/arrangement.py: 94.74%

1"""Best-effort, per-user storage of the GUI's panel and tab arrangement. 

2 

3Lives in the user's own cache directory (never in the shared project 

4folder), so one person's layout cannot overwrite another's. 

5""" 

6 

7import hashlib 

8import logging 

9import os 

10import re 

11import tempfile 

12from pathlib import Path 

13 

14import yaml 

15from platformdirs import user_cache_dir 

16 

17logger = logging.getLogger(__name__) 

18 

19VERSION = 1 

20TAB_KINDS = {"dag", "node"} 

21DIRECTIONS = {"left", "right", "above", "below"} 

22MAX_LAYOUT_OPS = 200 

23PIXEL_FLEX = re.compile(r"^0 [01] \d+(\.\d+)?(px|%)$") 

24 

25 

26def sanitize_arrangement(data): 

27 """Keep only the parts of a saved arrangement the GUI understands. 

28 

29 Returns None when ``data`` is not a usable arrangement. 

30 """ 

31 if not isinstance(data, dict) or data.get("version") != VERSION: 

32 return None 

33 result = {"version": VERSION} 

34 if isinstance(data.get("focused"), str): 

35 result["focused"] = data["focused"] 

36 panels = {} 

37 raw_panels = data.get("panels") 

38 for name, panel in (raw_panels.items() if isinstance(raw_panels, dict) else []): 

39 if not isinstance(name, str) or not isinstance(panel, dict): 

40 continue 

41 tabs = [tab for tab in map(_sanitize_tab, panel.get("tabs") or []) if tab] 

42 active = panel.get("active") 

43 if type(active) is not int or not 0 <= active < len(tabs): 

44 active = 0 

45 panels[name] = {"active": active, "tabs": tabs} 

46 result["panels"] = panels 

47 # Absent in saves from before this was remembered: then nothing is changed on restore. 

48 if isinstance(data.get("collapsed"), list): 

49 result["collapsed"] = [ 

50 name for name in data["collapsed"] if isinstance(name, str) and name 

51 ] 

52 result["layout"] = _sanitize_layout(data.get("layout")) 

53 result["sizes"] = _sanitize_sizes(data.get("sizes")) 

54 return result 

55 

56 

57def _sanitize_layout(ops): 

58 """The log of panel splits and closes, replayed in order on restore.""" 

59 clean = [] 

60 for op in ops if isinstance(ops, list) else []: 

61 if not isinstance(op, dict) or not isinstance(op.get("key"), str): 

62 continue 

63 if op.get("op") == "split": 

64 if isinstance(op.get("from"), str) and op.get("direction") in DIRECTIONS: 

65 clean.append( 

66 { 

67 "op": "split", 

68 "from": op["from"], 

69 "direction": op["direction"], 

70 "key": op["key"], 

71 } 

72 ) 

73 elif op.get("op") == "remove": 

74 clean.append({"op": "remove", "key": op["key"]}) 

75 return clean[-MAX_LAYOUT_OPS:] 

76 

77 

78def _sanitize_sizes(sizes): 

79 return { 

80 name: size 

81 for name, size in (sizes.items() if isinstance(sizes, dict) else []) 

82 if isinstance(name, str) and isinstance(size, str) and PIXEL_FLEX.match(size) 

83 } 

84 

85 

86def _sanitize_tab(tab): 

87 if not isinstance(tab, dict) or tab.get("kind") not in TAB_KINDS: 

88 return None 

89 if tab["kind"] == "dag": 

90 return {"kind": "dag"} 

91 group = tab.get("group") 

92 if not isinstance(group, str) or not group: 

93 return None 

94 return {"kind": "node", "group": group} 

95 

96 

97class ArrangementStore: 

98 def __init__( 

99 self, 

100 pipeline_directory: Path, 

101 pipeline_name: str | None = None, 

102 default_layout: Path | None = None, 

103 ): 

104 key = hashlib.sha256( 

105 str(Path(pipeline_directory).resolve()).encode() 

106 ).hexdigest() 

107 directory = ( 

108 Path(user_cache_dir("collectra", appauthor=False)) 

109 / "gui-arrangements" 

110 ) 

111 name = pipeline_name or Path(pipeline_directory).name 

112 name = re.sub(r'[/\\\x00-\x1f]', "_", name).strip() 

113 if name in {"", ".", ".."}: 

114 name = "pipeline" 

115 self.path = directory / name / f"{key}.yaml" 

116 self.legacy_path = directory / f"{key}.yaml" 

117 self.default_layout = default_layout 

118 

119 def read(self): 

120 for path in (self.path, self.legacy_path, self.default_layout): 

121 if path is not None: 

122 arrangement = self._read_file(path) 

123 if arrangement is not None: 

124 return arrangement 

125 return None 

126 

127 @staticmethod 

128 def _read_file(path): 

129 try: 

130 return sanitize_arrangement( 

131 yaml.safe_load(path.read_text(encoding="utf-8")) 

132 ) 

133 except (OSError, UnicodeError, yaml.YAMLError): 

134 return None 

135 

136 def custom_layouts(self): 

137 try: 

138 return sorted( 

139 (p for p in self.path.parent.glob("*.y*ml") 

140 if p != self.path and not re.fullmatch(r"[0-9a-f]{64}", p.stem) 

141 and p.suffix in {".yaml", ".yml"}), 

142 key=lambda p: p.name.casefold(), 

143 ) 

144 except OSError: 

145 return [] 

146 

147 def save_named(self, name: str, arrangement: dict): 

148 if (not isinstance(name, str) or not name.strip() 

149 or name.strip() in {".", ".."} 

150 or re.search(r'[/\\\x00-\x1f]', name) 

151 or re.fullmatch(r"[0-9a-f]{64}", name.strip())): 

152 raise ValueError("Enter a layout name without slashes or control characters.") 

153 path = self.path.parent / (name.strip() + ".yaml") 

154 if path.exists(): 

155 raise ValueError("A layout with that name already exists. Choose another name.") 

156 self._write_file(path, arrangement, raise_errors=True) 

157 

158 def write(self, arrangement: dict) -> None: 

159 self._write_file(self.path, arrangement) 

160 

161 @staticmethod 

162 def _write_file(path, arrangement, raise_errors=False): 

163 temporary = None 

164 try: 

165 path.parent.mkdir(parents=True, exist_ok=True) 

166 with tempfile.NamedTemporaryFile( 

167 mode="w", encoding="utf-8", dir=path.parent, delete=False 

168 ) as stream: 

169 temporary = Path(stream.name) 

170 yaml.safe_dump(arrangement, stream, sort_keys=False) 

171 os.replace(temporary, path) 

172 except OSError as error: 

173 logger.warning("Could not save GUI arrangement: %s", error) 

174 if raise_errors: 

175 raise 

176 finally: 

177 if temporary is not None: 

178 try: 

179 temporary.unlink(missing_ok=True) 

180 except OSError: 

181 pass