common.py (1412B)
1 from typing import Any, Callable 2 import json 3 import pathlib 4 import tarfile 5 import zipfile 6 7 import tqdm 8 9 10 def map_minhash(minhash_path: pathlib.Path, function: Callable[[dict[str, Any]], None]) -> None: 11 if minhash_path.name.endswith('.zip'): 12 return map_minhash_zip(minhash_path, function) 13 if minhash_path.name.endswith('.tar.gz') or minhash_path.name.endswith('.tgz') : 14 return map_minhash_tgz(minhash_path, function) 15 raise NotImplemented("Minhash archive format.") 16 17 18 def map_minhash_tgz(minhash_path: pathlib.Path, function: Callable[[dict[str, Any]], None]) -> None: 19 with tarfile.open(minhash_path, "r:gz") as tar: 20 for fileobject in tqdm.tqdm(tar.getmembers(), desc="Loading minhash output"): 21 if not fileobject.name.endswith(".out"): 22 continue 23 file = tar.extractfile(fileobject) 24 if file is not None: 25 for line in file: 26 function(json.loads(line)) 27 28 29 def map_minhash_zip(minhash_path: pathlib.Path, function: Callable[[dict[str, Any]], None]) -> None: 30 with zipfile.ZipFile(minhash_path) as zfile: 31 for filename in tqdm.tqdm(zfile.namelist(), desc="Loading minhash output"): 32 if not filename.endswith(".out"): 33 continue 34 with zfile.open(filename, "r") as file: 35 for line in file: 36 function(json.loads(line))