async def test_writer(compression: str):
baby_name_filename = os.path.join(os.path.dirname(__file__), 'data', 'baby_names.csv')
with NamedTemporaryFile(mode='wb', delete=False) as tmpfd:
async with aiofiles.open(baby_name_filename, 'rb') as rfd:
async with aiofiles.open(tmpfd.name, 'wb') as owfd:
async with asyncstream.open(owfd, mode='wb', compression=compression) as wfd:
async for line in rfd:
await wfd.write(line)
assert get_raw_lines(baby_name_filename) == get_raw_lines(tmpfd.name, compression)
async def test_open_read_with_filename(compression: str):
baby_name_filename = os.path.join(os.path.dirname(__file__), 'data',
'baby_names.csv')
with NamedTemporaryFile() as tmpfd:
async with aiofiles.open(baby_name_filename, 'rb') as cfd:
async with asyncstream.open(tmpfd.name,
mode='wb',
compression=compression) as fd:
await fd.write(await cfd.read())
assert get_raw_lines(baby_name_filename) == get_raw_lines(
tmpfd.name, compression)
async def test_open_read_parquet_with_filename(compression: str):
baby_name_filename = os.path.join(os.path.dirname(__file__), 'data', 'baby_names.csv')
with NamedTemporaryFile(delete=False) as tmpfd:
tmpfd.write(encode_parquet(baby_name_filename, compression))
tmpfd.flush()
async with asyncstream.open(tmpfd.name, mode='rt', encoding='parquet') as fd:
assert get_raw_lines(baby_name_filename) == await async_gen_to_list(fd)
async def test_open_read(compression: str):
baby_name_filename = os.path.join(os.path.dirname(__file__), 'data', 'baby_names.csv')
with NamedTemporaryFile() as tmpfd:
tmpfd.write(compress(baby_name_filename, compression))
tmpfd.flush()
async with aiofiles.open(tmpfd.name, 'rb') as cfd:
async with asyncstream.open(cfd, mode='rt', compression=compression) as fd:
assert get_raw_lines(baby_name_filename) == await async_gen_to_list(fd)