105 lines
3.7 KiB
Python
105 lines
3.7 KiB
Python
"""Implementations behind the CLI commands: parse, split, sort, merge, export."""
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from tqdm import tqdm
|
|
|
|
from collectionparser.utils import create_dir_for_file, find_files
|
|
|
|
# ponytail: buckets under this size sort byte-exactly in RAM; larger buckets
|
|
# shell out to `sort -u`, which spills to disk.
|
|
_IN_MEMORY_SORT_LIMIT = 8 * 1024 * 1024
|
|
|
|
|
|
class API:
|
|
"""Command implementations shared by the CLI."""
|
|
|
|
@staticmethod
|
|
def merge(src: str, dest: str) -> None:
|
|
"""Merge the bucket tree ``src`` into ``dest`` (destructive)."""
|
|
src_root = Path(src).expanduser()
|
|
dest_root = Path(dest).expanduser()
|
|
dest_files = {
|
|
str(f).replace(str(dest_root), ""): f for f in find_files(dest_root)
|
|
}
|
|
for src_file in find_files(src_root):
|
|
src_index = str(src_file).replace(str(src_root), "")
|
|
if src_index in dest_files:
|
|
append_file(src_file, dest_files[src_index])
|
|
else:
|
|
dst = dest_root / src_index.lstrip("/")
|
|
create_dir_for_file(str(dst))
|
|
move_file(src_file, dst)
|
|
|
|
@staticmethod
|
|
def sort(path: str) -> None:
|
|
"""Sort and de-duplicate every file under ``path`` in place."""
|
|
files = [p for p in find_files(path) if p.suffix != ".sorted"]
|
|
for file in tqdm(files, desc="Unique > Sorting > Replacing existing files"):
|
|
output_path = file.with_suffix(".sorted")
|
|
sort_unique(file, output_path)
|
|
shutil.move(output_path, file)
|
|
|
|
@staticmethod
|
|
def export_jsonl(dataset: str, out: str) -> None:
|
|
"""Export the dataset as JSON Lines (`{"email":..., "password":...}`)."""
|
|
create_dir_for_file(out)
|
|
total = 0
|
|
with open(out, "w", encoding="utf8") as writer:
|
|
for path in tqdm(find_files(dataset), desc="export jsonl"):
|
|
with open(path, encoding="utf8", errors="ignore") as reader:
|
|
for line in reader:
|
|
email, _, password = line.rstrip("\n").partition(":")
|
|
writer.write(
|
|
json.dumps(
|
|
{"email": email, "password": password},
|
|
ensure_ascii=False,
|
|
)
|
|
+ "\n"
|
|
)
|
|
total += 1
|
|
print(f"Exported {total:,} records to {out}.")
|
|
|
|
|
|
def sort_unique(src, dst) -> None:
|
|
"""Write the sorted, de-duplicated lines of ``src`` to ``dst``.
|
|
|
|
In-memory for small files, ``sort -u`` for large ones. ``src`` may equal
|
|
``dst`` (in-place).
|
|
"""
|
|
src = Path(src)
|
|
if src.stat().st_size <= _IN_MEMORY_SORT_LIMIT:
|
|
lines = src.read_bytes().split(b"\n")
|
|
if lines and lines[-1] == b"":
|
|
lines.pop()
|
|
lines = sorted(set(lines))
|
|
Path(dst).write_bytes(b"\n".join(lines) + b"\n" if lines else b"")
|
|
else:
|
|
tmp = f"{dst}.tmp"
|
|
with open(tmp, "wb") as writer:
|
|
subprocess.run(
|
|
["sort", "-u", str(src)],
|
|
stdout=writer,
|
|
check=True,
|
|
env={**os.environ, "LC_ALL": "C"},
|
|
)
|
|
os.replace(tmp, dst)
|
|
|
|
|
|
def append_file(src, dst) -> None:
|
|
"""Append ``src`` to ``dst`` and delete ``src`` (streams, constant memory)."""
|
|
print(f"merge {src} {dst}")
|
|
with open(src, "rb") as reader, open(dst, "ab") as writer:
|
|
shutil.copyfileobj(reader, writer, length=1024 * 1024)
|
|
print(f"del {src}")
|
|
os.remove(src)
|
|
|
|
|
|
def move_file(src, dst) -> None:
|
|
print(f"mv {src} {dst}")
|
|
shutil.move(src, dst)
|