def test_simple_sans_gz(self): sink = DiskSink("coba/tests/.temp/test.log") sink.write("a") sink.write("b") sink.write("c") self.assertEqual( ["a", "b", "c"], Path("coba/tests/.temp/test.log").read_text().splitlines())
def test_simple_with_gz(self): sink = DiskSink("coba/tests/.temp/test.gz") sink.write("a") sink.write("b") sink.write("c") lines = gzip.decompress( Path("coba/tests/.temp/test.gz").read_bytes()).decode( 'utf-8').splitlines() self.assertEqual(["a", "b", "c"], lines)
class TransactionIO_V4(Source['Result'], Sink[Any]): def __init__(self, log_file: Optional[str] = None, minify:bool=True) -> None: self._log_file = log_file self._minify = minify self._source = DiskSource(log_file) if log_file else ListSource() self._sink = DiskSink(log_file) if log_file else ListSink(self._source.items) def write(self, item: Any) -> None: if isinstance(self._sink, ListSink): self._sink.write(self._encode(item)) else: if not Path(self._sink._filename).exists():self._sink.write('["version",4]') self._sink.write(JsonEncode(self._minify).filter(self._encode(item))) def read(self) -> 'Result': n_lrns = None n_sims = None lrn_rows = {} env_rows = {} int_rows = {} if isinstance(self._source, ListSource): decoded_source = self._source else: decoded_source = Pipes.join(self._source, Foreach(JsonDecode())) for trx in decoded_source.read(): if not trx: continue if trx[0] == "experiment": n_lrns = trx[1]["n_learners"] n_sims = trx[1]["n_environments"] if trx[0] == "E": env_rows[trx[1]] = trx[2] if trx[0] == "L": lrn_rows[trx[1]] = trx[2] if trx[0] == "I": int_rows[tuple(trx[1])] = trx[2] return Result(n_lrns, n_sims, env_rows, lrn_rows, int_rows) def _encode(self,item): if item[0] == "T0": return ['experiment', {"n_learners":item[1], "n_environments":item[2]}] if item[0] == "T1": return ["L", item[1], item[2]] if item[0] == "T2": return ["E", item[1], item[2]] if item[0] == "T3": rows_T = collections.defaultdict(list) for row in item[2]: for col,val in row.items(): if col == "rewards" : col="reward" if col == "reveals" : col="reveal" rows_T[col].append(val) return ["I", item[1], { "_packed": rows_T }] return None