Coverage for collectra/commons/files.py: 95.56%

1import uuid 

2from datetime import datetime 

3from pathlib import Path 

4 

5import pydantic 

6 

7from ..types.base import load_artefact 

8from ..utils import write_yaml 

9 

10 

11class CollectraResultsMetadata(pydantic.BaseModel): 

12 # Some results.yaml files have an unquoted version (e.g. `version: 1.0`), 

13 # which PyYAML parses as a float rather than a str. 

14 model_config = pydantic.ConfigDict(coerce_numbers_to_str=True) 

15 

16 workflow: str | None = None 

17 version: str | None = None 

18 partition: str | None = None 

19 timestamp: str | None = None 

20 

21 def model_dump(self, *args, **kwargs) -> dict: 

22 data = super().model_dump(*args, **kwargs) 

23 data["timestamp"] = datetime.now().isoformat() 

24 if self.partition is None: 

25 data.pop("partition") 

26 if self.workflow is None: 

27 data.pop("workflow") 

28 if self.version is None: 

29 data.pop("version") 

30 return data 

31 

32 

33class CollectraFile(pydantic.BaseModel): 

34 collectra_file_path: Path 

35 collectra_results_metadata: CollectraResultsMetadata 

36 data: dict 

37 assets: dict[str, Path] = dict() 

38 cleanup_on_save_failure: bool = False 

39 

40 def __getitem__(self, key): 

41 item = self.data[key] 

42 return item 

43 

44 def __contains__(self, key): 

45 return key in self.data 

46 

47 @property 

48 def results_path(self) -> Path: 

49 """Return the canonical results file for this Collectra directory.""" 

50 return self.collectra_file_path / "results.yaml" 

51 

52 def model_dump(self, *args, **kwargs) -> dict: 

53 data = super().model_dump(*args, **kwargs) 

54 data["collectra_results_metadata"] = self.collectra_results_metadata.model_dump( 

55 *args, **kwargs 

56 ) 

57 for key, value in data["data"].items(): 

58 data[key] = ( 

59 value 

60 if not isinstance(value, list) 

61 else value[0] if isinstance(value, list) and len(value) == 1 else value 

62 ) 

63 data.pop("data") 

64 data.pop("assets") 

65 data.pop("cleanup_on_save_failure") 

66 data.pop("collectra_file_path") 

67 return data 

68 

69 @classmethod 

70 def from_data(cls, collectra_file_path: Path | str): 

71 import yaml 

72 

73 collectra_file_path = Path(collectra_file_path) 

74 results_yaml = collectra_file_path / "results.yaml" 

75 if not results_yaml.exists(): 

76 raise FileNotFoundError(f"results.yaml not found in {collectra_file_path}") 

77 with open(results_yaml, "r") as f: 

78 data = yaml.safe_load(f) or {} 

79 if not isinstance(data, dict): 

80 raise ValueError(f"Invalid results data in {results_yaml}") 

81 metadata = CollectraResultsMetadata( 

82 **data.pop("collectra_results_metadata", {}) 

83 ) 

84 return cls( 

85 collectra_file_path=collectra_file_path, 

86 collectra_results_metadata=metadata, 

87 data=data, 

88 ) 

89 

90 @classmethod 

91 def from_file( 

92 cls, 

93 file: Path, 

94 label: str, 

95 ext: str, 

96 output: Path | None = None, 

97 force: bool = False, 

98 **kwargs, 

99 ): 

100 ext = ext.lstrip(".") 

101 collectra_file_path = file.with_suffix(f".{ext}") 

102 if output: 

103 collectra_file_path = output / collectra_file_path.name 

104 directory_existed = collectra_file_path.exists() 

105 if directory_existed and not force: 

106 raise FileExistsError( 

107 f"File {collectra_file_path} already exists. To overwrite, use -f or --force option." 

108 ) 

109 else: 

110 collectra_file_path.mkdir(parents=True, exist_ok=True) 

111 partition = kwargs.get("partition", None) 

112 metadata = CollectraResultsMetadata(partition=partition) 

113 new_data = { 

114 f"{label}": { 

115 "id": f"{label}", 

116 "type": "collectra.Image", 

117 "data": str(file.name), 

118 } 

119 } 

120 return cls( 

121 collectra_file_path=collectra_file_path, 

122 collectra_results_metadata=metadata, 

123 data=new_data, 

124 assets={file.name: Path(file)}, 

125 cleanup_on_save_failure=not directory_existed, 

126 ) 

127 

128 def save(self): 

129 import copy 

130 import shutil 

131 

132 collectra_file_path = self.collectra_file_path 

133 assets = copy.deepcopy(self.assets) 

134 data = self.model_dump() 

135 try: 

136 for asset_name, asset_path in assets.items(): 

137 if not asset_path.exists(): 

138 raise FileNotFoundError(f"Asset file {asset_path} not found.") 

139 shutil.copy(asset_path, collectra_file_path / asset_name) 

140 write_yaml(data, self.results_path) 

141 self.cleanup_on_save_failure = False 

142 except Exception as e: 

143 if self.cleanup_on_save_failure: 

144 shutil.rmtree(collectra_file_path, ignore_errors=True) 

145 raise RuntimeError(f"Failed to save Collectra file: {e}")