| Total Complexity | 4 |
| Total Lines | 35 |
| Duplicated Lines | 0 % |
| Changes | 0 | ||
| 1 | from typing import Sequence |
||
|
|
|||
| 2 | |||
| 3 | import pandas as pd |
||
| 4 | |||
| 5 | from mandos.model.hits import AbstractHit, HitFrame |
||
| 6 | from mandos.model.concrete_hits import HIT_CLASSES |
||
| 7 | |||
| 8 | |||
| 9 | class HitUtils: |
||
| 10 | @classmethod |
||
| 11 | def hits_to_df(cls, hits: Sequence[AbstractHit]) -> HitFrame: |
||
| 12 | data = [] |
||
| 13 | for hit in hits: |
||
| 14 | x = {f: getattr(hit, f) for f in hit.__class__.fields()} |
||
| 15 | x["universal_id"] = hit.universal_id |
||
| 16 | x["hit_class"] = hit.hit_class |
||
| 17 | data.append(x) |
||
| 18 | return HitFrame([pd.Series(x) for x in data]) |
||
| 19 | |||
| 20 | @classmethod |
||
| 21 | def df_to_hits(cls, self: HitFrame) -> Sequence[AbstractHit]: |
||
| 22 | hits = [] |
||
| 23 | for row in self.itertuples(): |
||
| 24 | clazz = HIT_CLASSES[row.hit_class] |
||
| 25 | # ignore extra columns |
||
| 26 | # if cols are missing, let it fail on clazz.__init__ |
||
| 27 | data = {f: getattr(row, f) for f in clazz.fields()} |
||
| 28 | # noinspection PyArgumentList |
||
| 29 | hit = clazz(**data) |
||
| 30 | hits.append(hit) |
||
| 31 | return hits |
||
| 32 | |||
| 33 | |||
| 34 | __all__ = ["HitUtils"] |
||
| 35 |