| 179 | |
| 180 | |
| 181 | class ResultCollection: |
| 182 | def __init__(self, *results): |
| 183 | self.results = {result.primary_key: result for result in results} |
| 184 | |
| 185 | @property |
| 186 | def primary_key_names(self): |
| 187 | return next(iter(self.results.values())).primary_key_names |
| 188 | |
| 189 | def toframe(self) -> pd.DataFrame: |
| 190 | """Convert results to pandas dataframe.""" |
| 191 | return pd.DataFrame([result.todict() for result in self.results.values()]).set_index( |
| 192 | list(self.primary_key_names) |
| 193 | ) |
| 194 | |
| 195 | def add(self, result: Result): |
| 196 | """Add a result to the collection.""" |
| 197 | if result.primary_key in self.results: |
| 198 | raise ValueError( |
| 199 | "An entry for {result.primary_key} does already " |
| 200 | "exist in this collection. Did you try to add the " |
| 201 | "same result twice?" |
| 202 | ) |
| 203 | if len(self) > 0: |
| 204 | if result.primary_key_names != self.primary_key_names: |
| 205 | raise ValueError("Incompatible result format.") |
| 206 | self.results[result.primary_key] = result |
| 207 | |
| 208 | @classmethod |
| 209 | def fromdicts(cls, data: Iterable[dict]): |
| 210 | return cls(*[Result.fromdict(entry) for entry in data]) |
| 211 | |
| 212 | def todicts(self): |
| 213 | return [result.todict() for result in self.results.values()] |
| 214 | |
| 215 | def __len__(self): |
| 216 | return len(self.results) |
| 217 | |
| 218 | def __contains__(self, other: Result): |
| 219 | if not isinstance(other, Result): |
| 220 | raise ValueError(f"{type(self)} can only store objects of type Result, but got {type(other)}.") |
| 221 | return other.primary_key in self.results |
| 222 | |
| 223 | def __eq__(self, other): |
| 224 | if not isinstance(other, ResultCollection): |
| 225 | return False |
| 226 | return other.results == self.results |