diff --git a/dataforge/sources.py b/dataforge/sources.py new file mode 100644 index 0000000..2e91809 --- /dev/null +++ b/dataforge/sources.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import csv as _csv +import json +import os +import time +from typing import Optional + +from .core import Table +from .utils import norm_path, parse_partition_filter + +BOM_ENCODING = "utf-8-sig" + + +def _read_csv_rows(path, thousands_sep: Optional[str] = None, dialect=None): + with open(norm_path(path), "r", encoding=BOM_ENCODING, newline="") as f: + reader = _csv.DictReader(f, dialect=dialect) + for raw in reader: + row = {} + for k, v in raw.items(): + key = k.lstrip("\ufeff") + if thousands_sep and v and "," in v: + try: + v = v.replace(thousands_sep, "") + except ValueError: + pass + row[key] = v + yield row + + +def csv(path, thousands_sep: Optional[str] = None, partition_filter: Optional[str] = None, + dialect=None) -> Table: + """读取单个文件或 glob 模式(glob 支持见 PR #55)。 + + partition_filter 见 discussion #5 的用法示例。 + """ + import glob as _glob + paths = [path] + if any(ch in path for ch in "*?["): + paths = sorted(_glob.glob(norm_path(path))) + if not paths: + raise FileNotFoundError(path) + if partition_filter: + col, op, value = parse_partition_filter(partition_filter) + paths = [p for p in paths if _match(p, col, op, value)] + rows = [] + for p in paths: + rows.extend(_read_csv_rows(p, thousands_sep=thousands_sep, dialect=dialect)) + return Table(rows) + + +def _match(path, col, op, value): + # 目录名里找 dt=YYYY-MM-DD 这类分区键 + import re + m = re.search(rf"{col}=([^/]+)", path) + if not m: + return True + return _cmp(m.group(1), op, value) + + +def _cmp(a, op, b): + try: + a, b = float(a), float(b) + except ValueError: + pass + if op == ">=": + return a >= b + if op == "<=": + return a <= b + if op == "==": + return a == b + if op == "!=": + return a != b + if op == ">": + return a > b + if op == "<": + return a < b + raise ValueError(op) + + +def jsonl(path) -> Table: + rows = [] + with open(norm_path(path), "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + return Table(rows) + + +class _HTTPReader: + """HTTP 源:带退避重试(上限 30s + full jitter,PR #65)。""" + MAX_BACKOFF = 30.0 + + def __init__(self, url, timeout=5.0, retries=4): + self.url = url + self.timeout = timeout + self.retries = retries + + def read(self): + import random + import urllib.error + import urllib.request + delay = 1.0 + last = None + for attempt in range(self.retries): + try: + with urllib.request.urlopen(self.url, timeout=self.timeout) as r: + return r.read() + except urllib.error.URLError as e: + last = e + time.sleep(min(delay, self.MAX_BACKOFF) * random.random() * 2) + delay *= 2 + raise last + + +def http(url, **kw) -> Table: + data = _HTTPReader(url, **kw).read() + return jsonl_from_bytes(data) + + +def jsonl_from_bytes(data) -> Table: + import io + rows = [] + for line in io.StringIO(data.decode("utf-8")): + line = line.strip() + if line: + rows.append(json.loads(line)) + return Table(rows) + + +def s3(bucket_prefix, **kw): + """s3 源(可选依赖 boto3)。分区目录逐文件读取。""" + try: + import boto3 # noqa: F401 + except ImportError as e: # pragma: no cover + raise RuntimeError("pip install dataforge[s3]") from e + raise NotImplementedError("s3 源需要 dataforge[s3] 与凭证;本地用 csv() 即可") + + +def read(cfg=None): + """connect() 的底层入口。""" + return csv(cfg.get("path"), **cfg.get("csv_opts", {}))