24 lines
873 B
Python
24 lines
873 B
Python
import os
|
||
import tempfile
|
||
import zipfile
|
||
|
||
from gnnrec.kgrec.utils import iter_json
|
||
|
||
|
||
def iter_lines(raw_path, data_type):
|
||
"""依次迭代OAG数据集某种类型数据所有txt文件的每一行并将JSON解析为字典
|
||
|
||
:param raw_path: str 原始zip文件所在目录
|
||
:param data_type: str 数据类型，author, paper, venue, affiliation之一
|
||
:return: Iterable[dict]
|
||
"""
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
for zip_file in os.listdir(raw_path):
|
||
if zip_file.startswith(f'mag_{data_type}s'):
|
||
with zipfile.ZipFile(os.path.join(raw_path, zip_file)) as z:
|
||
for txt in z.namelist():
|
||
print(f'{zip_file}\\{txt}')
|
||
txt_file = z.extract(txt, tmp)
|
||
yield from iter_json(txt_file)
|
||
os.remove(txt_file)
|