diff --git a/lib/bgcalc/freqs/storage.py b/lib/bgcalc/freqs/storage.py index 88e0a86371..d03fb574b3 100644 --- a/lib/bgcalc/freqs/storage.py +++ b/lib/bgcalc/freqs/storage.py @@ -102,7 +102,7 @@ async def wrapper(args: FreqCalcArgs) -> FreqCalcResult: data, cache_path = await find_cached_result(args) if data is None: data: FreqCalcResult = await func(args) - async with aiofiles.open(cache_path, 'w') as fw, AsyncBatchWriter(fw, 100) as bw: + async with AsyncBatchWriter(cache_path, 'w', 100) as bw: common_md = CommonMetadata(num_blocks=len(data.freqs), conc_size=data.conc_size) await bw.write(json.dumps(common_md.to_dict()) + '\n') for freq in data.freqs: diff --git a/lib/bgcalc/keywords/__init__.py b/lib/bgcalc/keywords/__init__.py index f8f980ab61..30d7869b17 100644 --- a/lib/bgcalc/keywords/__init__.py +++ b/lib/bgcalc/keywords/__init__.py @@ -108,7 +108,7 @@ async def wrapper(corp: KCorpus, ref_corp: KCorpus, args: KeywordsFormArgs, max_ ans = await f(corp, ref_corp, args, max_items) # ans = sorted(ans, key=lambda x: x[1], reverse=True) num_lines = len(ans) - async with aiofiles.open(path, 'w') as fw, AsyncBatchWriter(fw, 100) as bw: + async with AsyncBatchWriter(path, 'w', 100) as bw: await bw.write(json.dumps(dict(total=num_lines)) + '\n') for item in ans: await bw.write(item.to_json() + '\n') diff --git a/lib/bgcalc/wordlist/__init__.py b/lib/bgcalc/wordlist/__init__.py index a4577c84f1..1a29aa7e7e 100644 --- a/lib/bgcalc/wordlist/__init__.py +++ b/lib/bgcalc/wordlist/__init__.py @@ -77,7 +77,7 @@ async def wrapper(corp: AbstractKCorpus, args: WordlistFormArgs, max_items: int) ans = await f(corp, args, sys.maxsize) ans = sorted(ans, key=lambda x: x[1], reverse=True) num_lines = len(ans) - async with aiofiles.open(path, 'w') as fw, AsyncBatchWriter(fw, 100) as bw: + async with AsyncBatchWriter(path, 'w', 100) as bw: await bw.write(json.dumps(dict(total=num_lines)) + '\n') for item in ans: await bw.write(json.dumps(item) + '\n') diff --git a/lib/plugins/default_subcmixer/__init__.py b/lib/plugins/default_subcmixer/__init__.py index 17b775412b..4050040ee9 100644 --- a/lib/plugins/default_subcmixer/__init__.py +++ b/lib/plugins/default_subcmixer/__init__.py @@ -20,7 +20,6 @@ from collections import defaultdict from typing import Any, Dict, List -import aiofiles import plugins import ujson as json from action.control import http_action @@ -75,7 +74,7 @@ async def subcmixer_create_subcorpus(amodel: CorpusActionModel, req: KRequest, r struct_indices = sorted([int(x) for x in req.form.get('ids').split(',')]) id_attr = req.form.get('idAttr').split('.') attr = amodel.corp.get_struct(id_attr[0]) - async with aiofiles.open(os.path.join(amodel.subcpath, subc_id.data_path), 'wb') as fw, AsyncBatchWriter(fw, 100) as bw: + async with AsyncBatchWriter(os.path.join(amodel.subcpath, subc_id.data_path), 'wb', 100) as bw: for idx in struct_indices: await bw.write(struct.pack(' bool: async def export_jsonl(data: List[DocListItem], target_path: str) -> bool: if len(data) == 0: return False - async with aiofiles.open(target_path, 'w') as fw, AsyncBatchWriter(fw, 100) as bw: + async with AsyncBatchWriter(target_path, 'w', 100) as bw: for item in data: await bw.write(json.dumps(item.to_dict()) + "\n") return True diff --git a/lib/plugins/masm_subcmixer/__init__.py b/lib/plugins/masm_subcmixer/__init__.py index 9734ca58a5..c13b391266 100644 --- a/lib/plugins/masm_subcmixer/__init__.py +++ b/lib/plugins/masm_subcmixer/__init__.py @@ -21,7 +21,6 @@ from dataclasses import dataclass from typing import List, Optional -import aiofiles import aiohttp import plugins import ujson as json @@ -79,7 +78,7 @@ async def subcmixer_create_subcorpus(amodel: CorpusActionModel, req: KRequest, r struct_idxs = sorted(attr.str2id(sid) for sid in struct_ids) subc_id = await create_new_subc_ident(amodel.subcpath, amodel.corp.corpname) - async with aiofiles.open(os.path.join(amodel.subcpath, subc_id.data_path), 'wb') as fw, AsyncBatchWriter(fw, 100) as bw: + async with AsyncBatchWriter(os.path.join(amodel.subcpath, subc_id.data_path), 'wb', 100) as bw: for idx in struct_idxs: await bw.write(struct.pack(' str: class AsyncBatchWriter(AsyncContextDecorator): - def __init__(self, f, batch_size: int): - self.f = f + def __init__(self, filename: str, mode: str, batch_size: int): + self.filename = filename + self.mode = mode + self.f = None self.batch_size = batch_size self.lines = [] @@ -65,8 +69,10 @@ async def flush(self): self.lines = [] async def __aenter__(self): + self.f = await aiofiles.open(self.filename, self.mode).__aenter__() return self async def __aexit__(self, *exc): await self.flush() + self.f.close() return False