| 23 | pass |
| 24 | |
| 25 | class DummyStorage(DataFlowStorage): |
| 26 | def __init__( |
| 27 | self, |
| 28 | cache_path:str=None, |
| 29 | file_name_prefix:str=None, |
| 30 | cache_type: Literal["json", "jsonl", "csv", "parquet", "pickle", None] = None |
| 31 | ): |
| 32 | self._data = None |
| 33 | self.cache_path = cache_path |
| 34 | self.file_name_prefix = file_name_prefix |
| 35 | self.cache_type = cache_type |
| 36 | |
| 37 | def set_data(self, data: Any): |
| 38 | """ |
| 39 | Set data to be written later. |
| 40 | """ |
| 41 | self._data = data |
| 42 | |
| 43 | def set_file_name_prefix(self, file_name_prefix: str): |
| 44 | """ |
| 45 | Set the file name prefix for cache files. |
| 46 | """ |
| 47 | self.file_name_prefix = file_name_prefix |
| 48 | |
| 49 | def read(self, output_type: Literal["dataframe", "dict"] = "dataframe") -> Any: |
| 50 | return self._data |
| 51 | |
| 52 | def write(self, data): |
| 53 | self._data = data |
| 54 | if self.cache_type != None: |
| 55 | cache_file_path = os.path.join(self.cache_path, f"{self.file_name_prefix}.{self.cache_type}") |
| 56 | os.makedirs(os.path.dirname(cache_file_path), exist_ok=True) |
| 57 | if self.cache_type == "json": |
| 58 | data.to_json(cache_file_path, orient="records", force_ascii=False, indent=2) |
| 59 | elif self.cache_type == "jsonl": |
| 60 | data.to_json(cache_file_path, orient="records", lines=True, force_ascii=False) |
| 61 | elif self.cache_type == "csv": |
| 62 | data.to_csv(cache_file_path, index=False) |
| 63 | elif self.cache_type == "parquet": |
| 64 | data.to_parquet(cache_file_path) |
| 65 | elif self.cache_type == "pickle": |
| 66 | data.to_pickle(cache_file_path) |
| 67 | else: |
| 68 | raise ValueError(f"Unsupported file type: {self.cache_type}, output file should end with json, jsonl, csv, parquet, pickle") |
| 69 | |
| 70 | class FileStorage(DataFlowStorage): |
| 71 | """ |