135 lines
3.7 KiB
Python
135 lines
3.7 KiB
Python
from typing import Any, TextIO, cast
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
import gzip, glob, json
|
|
|
|
class DatasetIO:
|
|
@staticmethod
|
|
def resolve_input_files(input_files: list[str], directory_pattern: str | tuple[str, ...] = ("*.jsonl", "*.jsonl.gz")) -> list[Path]:
|
|
patterns = (
|
|
(directory_pattern,)
|
|
if isinstance(directory_pattern, str)
|
|
else directory_pattern
|
|
)
|
|
resolved = []
|
|
seen = set()
|
|
|
|
for item in input_files:
|
|
path = Path(item)
|
|
if path.is_dir():
|
|
for pattern in patterns:
|
|
for file_path in sorted(path.rglob(pattern)):
|
|
key = str(file_path.resolve())
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
resolved.append(file_path)
|
|
continue
|
|
|
|
if any(ch in item for ch in "*?[]"):
|
|
for match in sorted(glob.glob(item)):
|
|
file_path = Path(match)
|
|
if not file_path.is_file():
|
|
continue
|
|
key = str(file_path.resolve())
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
resolved.append(file_path)
|
|
continue
|
|
|
|
if path.is_file():
|
|
key = str(path.resolve())
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
resolved.append(path)
|
|
|
|
return resolved
|
|
|
|
@staticmethod
|
|
def list_directory_files(input_dir:Path, directory_pattern:str) -> list[Path]:
|
|
if not input_dir.exists():
|
|
return []
|
|
return sorted(input_dir.rglob(directory_pattern))
|
|
|
|
@staticmethod
|
|
def open_text(path:Path, mode:str="r"):
|
|
text_mode = mode if "t" in mode else f"{mode}t"
|
|
open_fn = gzip.open if path.suffix == ".gz" else open
|
|
return cast(TextIO, open_fn(path, text_mode, encoding="utf-8"))
|
|
|
|
@staticmethod
|
|
def iter_jsonl_rows(path:Path):
|
|
with DatasetIO.open_text(path, "r") as handle:
|
|
for line in handle:
|
|
if line.strip():
|
|
yield json.loads(line)
|
|
|
|
@staticmethod
|
|
def append_jsonl_row(path:Path, row:dict[str, Any]):
|
|
with DatasetIO.open_text(path, "a") as raw_handle:
|
|
handle = cast(TextIO, raw_handle)
|
|
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
|
|
@staticmethod
|
|
def count_jsonl_rows(path:Path) -> int:
|
|
count = 0
|
|
|
|
candidates = [path]
|
|
if path.suffix == ".gz":
|
|
candidates.append(path.with_suffix(""))
|
|
else:
|
|
candidates.append(Path(f"{path}.gz"))
|
|
|
|
seen = set()
|
|
for candidate in candidates:
|
|
if candidate in seen:
|
|
continue
|
|
seen.add(candidate)
|
|
|
|
if not candidate.exists() or not candidate.is_file():
|
|
continue
|
|
|
|
try:
|
|
with DatasetIO.open_text(candidate, "r") as handle:
|
|
for line in handle:
|
|
if line.strip():
|
|
count += 1
|
|
except (OSError, UnicodeError):
|
|
continue
|
|
|
|
return count
|
|
|
|
@staticmethod
|
|
def rotate_and_gzip_if_size_reached(path:Path, max_bytes:int) -> bool:
|
|
if max_bytes <= 0:
|
|
return False
|
|
if not path.exists() or not path.is_file():
|
|
return False
|
|
if path.stat().st_size < max_bytes:
|
|
return False
|
|
|
|
timestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
|
|
if path.suffix == ".jsonl":
|
|
rotated_name = f"{path.stem}.{timestamp}.jsonl"
|
|
else:
|
|
rotated_name = f"{path.name}.{timestamp}"
|
|
|
|
rotated_path = path.with_name(rotated_name)
|
|
suffix = 1
|
|
while rotated_path.exists():
|
|
suffix += 1
|
|
if path.suffix == ".jsonl":
|
|
rotated_name = f"{path.stem}.{timestamp}.{suffix}.jsonl"
|
|
else:
|
|
rotated_name = f"{path.name}.{timestamp}.{suffix}"
|
|
rotated_path = path.with_name(rotated_name)
|
|
|
|
path.rename(rotated_path)
|
|
with rotated_path.open("rb") as src:
|
|
with gzip.open(f"{rotated_path}.gz", "wb") as dst:
|
|
dst.writelines(src)
|
|
rotated_path.unlink()
|
|
return True
|