quickquery 0.1.1__py3-none-any.whl

quickquery/live.py

@@ -0,0 +1,195 @@
+from contextlib import ExitStack
+from dataclasses import dataclass, fields
+from types import TracebackType
+from typing import Any, Self
+
+from camoufox.sync_api import Camoufox
+from patchright.sync_api import (
+    Page as PatchrightPage,
+    Playwright,
+    sync_playwright,
+)
+from playwright.sync_api import Page as PlaywrightPage
+
+Page = PatchrightPage | PlaywrightPage
+
+
+@dataclass(frozen=True, slots=True)
+class RecycleEvery:
+    browser: int | None = None
+    context: int | None = None
+    page: int | None = None
+
+    def __post_init__(self) -> None:
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if value is not None and value < 1:
+                raise ValueError(f'{f.name} は 1 以上で指定してください (got {value})')
+
+
+class _SessionBase:
+    def __init__(
+        self,
+        *,
+        browser_options: dict[str, Any] | None = None,
+        context_options: dict[str, Any] | None = None,
+        recycle: RecycleEvery | None = None,
+    ) -> None:
+        self._recycle = recycle or RecycleEvery()
+        self._browser_options = dict(browser_options or {})
+        self._context_options = dict(context_options or {})
+        self._browser = None
+        self._context = None
+        self._page: Page | None = None
+        self._page_calls = 0
+        self._entered = False
+
+    def page(self) -> Page:
+        if not self._entered:
+            raise RuntimeError('with ブロックの外で page() を呼べません')
+        if self._page is None:
+            self._open_browser()
+        elif (b := self._recycle.browser) and self._page_calls % b == 0:
+            self._close_browser()
+            self._open_browser()
+        elif (c := self._recycle.context) and self._page_calls % c == 0:
+            self._close_context()
+            self._open_context()
+        elif (p := self._recycle.page) and self._page_calls % p == 0:
+            self._close_page()
+            self._open_page()
+        self._page_calls += 1
+        return self._page
+
+    def _open_page(self) -> None:
+        self._page = self._context.new_page()
+
+    def _close_page(self) -> None:
+        if self._page is not None:
+            self._page.close()
+            self._page = None
+
+    def _open_context(self) -> None:
+        self._context = self._browser.new_context(**self._context_options)
+        self._open_page()
+
+    def _close_context(self) -> None:
+        self._close_page()
+        if self._context is not None:
+            self._context.close()
+            self._context = None
+
+
+class PatchrightSession(_SessionBase):
+    def __init__(
+        self,
+        *,
+        browser_options: dict[str, Any] | None = None,
+        context_options: dict[str, Any] | None = None,
+        recycle: RecycleEvery | None = None,
+    ) -> None:
+        super().__init__(
+            browser_options=browser_options,
+            context_options=context_options,
+            recycle=recycle,
+        )
+        self._pw: Playwright | None = None
+
+    def __enter__(self) -> Self:
+        self._pw = sync_playwright().start()
+        self._entered = True
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if not self._entered:
+            return
+        self._close_browser()
+        self._pw.stop()
+        self._pw = None
+        self._entered = False
+        self._page_calls = 0
+
+    def _open_browser(self) -> None:
+        self._browser = self._pw.chromium.launch(**self._browser_options)
+        self._open_context()
+
+    def _close_browser(self) -> None:
+        self._close_context()
+        if self._browser is not None:
+            self._browser.close()
+        self._browser = None
+
+
+class CamoufoxSession(_SessionBase):
+    def __init__(
+        self,
+        *,
+        browser_options: dict[str, Any] | None = None,
+        context_options: dict[str, Any] | None = None,
+        recycle: RecycleEvery | None = None,
+    ) -> None:
+        super().__init__(
+            browser_options=browser_options,
+            context_options=context_options,
+            recycle=recycle,
+        )
+        self._stack: ExitStack | None = None
+
+    def __enter__(self) -> Self:
+        self._entered = True
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if not self._entered:
+            return
+        self._close_browser()
+        self._entered = False
+        self._page_calls = 0
+
+    def _open_browser(self) -> None:
+        self._stack = ExitStack()
+        self._browser = self._stack.enter_context(Camoufox(**self._browser_options))
+        self._open_context()
+
+    def _close_browser(self) -> None:
+        self._close_context()
+        if self._stack is not None:
+            self._stack.close()
+            self._stack = None
+        self._browser = None
+
+
+def open_patchright(
+    *,
+    browser_options: dict[str, Any] | None = None,
+    context_options: dict[str, Any] | None = None,
+    recycle: RecycleEvery | None = None,
+) -> PatchrightSession:
+    return PatchrightSession(
+        browser_options=browser_options,
+        context_options=context_options,
+        recycle=recycle,
+    )
+
+
+def open_camoufox(
+    *,
+    browser_options: dict[str, Any] | None = None,
+    context_options: dict[str, Any] | None = None,
+    recycle: RecycleEvery | None = None,
+) -> CamoufoxSession:
+    return CamoufoxSession(
+        browser_options=browser_options,
+        context_options=context_options,
+        recycle=recycle,
+    )

quickquery/utils.py

@@ -0,0 +1,199 @@
+import csv
+import hashlib
+import html
+import os
+from collections.abc import Mapping
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+from typing import Callable, Iterable, Iterator
+
+from loguru import logger
+from selectolax.lexbor import LexborHTMLParser
+from tqdm import tqdm
+
+
+def _ensure_parent(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+
+def parse_html(path: Path) -> LexborHTMLParser | None:
+    try:
+        return LexborHTMLParser(path.read_bytes())
+    except Exception as e:
+        logger.error(f'[parse_html] {path} {type(e).__name__}: {e}')
+        return None
+
+
+def meta_html(meta: Mapping[str, object | None]) -> str:
+    return ''.join(
+        f'<meta name="{html.escape(name)}" content="{html.escape(str(content))}">'
+        for name, content in meta.items()
+        if content is not None
+    )
+
+
+def from_here(file: str) -> Callable[[str], Path]:
+    base = Path(file).resolve().parent
+    return lambda path: base / path
+
+def append_csv(path: Path, row: dict) -> None:
+    '''``row`` を 1 行だけ CSV に追記する（ファイルが無ければ作成）。
+
+    Excel 互換のため、**ファイル新規作成時のみ先頭に UTF-8 BOM** を書く
+    （``utf-8-sig`` で open）。既存ファイルへの追記では BOM を書かない
+    （中途 BOM は不正になるため）。ファイルが新規 / 空ならヘッダ行を書く。
+    列順は ``row.keys()`` の順で、2 回目以降のキーずれは検知しない
+    （pandas 版と同じ挙動）。
+    '''
+    try:
+        _ensure_parent(path)
+        need_header = not path.exists() or path.stat().st_size == 0
+        encoding = 'utf-8-sig' if need_header else 'utf-8'
+        with open(path, mode='a', newline='', encoding=encoding) as f:
+            w = csv.DictWriter(f, fieldnames=list(row.keys()))
+            if need_header:
+                w.writeheader()
+            w.writerow(row)
+    except Exception as e:
+        logger.error(f'[append_csv] {path} {row} {type(e).__name__}: {e}')
+
+def write_csv(path: Path, rows: list[dict]) -> None:
+    '''``rows`` を CSV ファイルとして書き出す（上書き）。
+
+    Excel 互換のため UTF-8 BOM（``utf-8-sig``）とヘッダ行を付ける。
+    ``rows`` が空ならスキップ（警告のみ）。列順は先頭行の ``keys()`` の順で、
+    2 回目以降のキーずれは検知しない（``append_csv`` と同じ）。
+    '''
+    try:
+        if not rows:
+            logger.warning(f'[write_csv] {path} no rows, skipped')
+            return
+        _ensure_parent(path)
+        with open(path, mode='w', newline='', encoding='utf-8-sig') as f:
+            w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+            w.writeheader()
+            w.writerows(rows)
+    except Exception as e:
+        logger.error(f'[write_csv] {path} {type(e).__name__}: {e}')
+
+def write_parquet(path: Path, rows: list[dict]) -> None:
+    '''``rows`` を Parquet ファイルとして書き出す。
+
+    pyarrow を直接使う（pandas 非依存）。``rows`` が空ならスキップ（警告のみ）。
+    列スキーマは各列の最初の non-None 値から推論されるので、**同一キーで型が
+    混在するとエラーになる**ことがある点に注意。
+    '''
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    try:
+        if not rows:
+            logger.warning(f'[write_parquet] {path} no rows, skipped')
+            return
+        _ensure_parent(path)
+        pq.write_table(pa.Table.from_pylist(rows), path)
+    except Exception as e:
+        logger.error(f'[write_parquet] {path} {type(e).__name__}: {e}')
+
+def hash_name(key: str) -> str:
+    return hashlib.md5(key.encode()).hexdigest()
+
+def write_text(path: Path, data: str) -> bool:
+    try:
+        _ensure_parent(path)
+        path.write_text(data, encoding='utf-8', errors='replace')
+        return True
+    except Exception as e:
+        logger.error(f'[write_text] {path} {type(e).__name__}: {e}')
+        return False
+
+def write_bytes(path: Path, data: bytes) -> bool:
+    try:
+        _ensure_parent(path)
+        path.write_bytes(data)
+        return True
+    except Exception as e:
+        logger.error(f'[write_bytes] {path} {type(e).__name__}: {e}')
+        return False
+
+def save_log(path: Path, level: str = 'WARNING') -> None:
+    '''コンソール（stderr）に出るログと同じ内容を、指定ファイルにも残す。'''
+    _ensure_parent(path)
+    logger.add(path, level=level, encoding='utf-8')
+
+
+class _SafeWorker:
+    def __init__(self, fn: Callable) -> None:
+        self.fn = fn
+
+    def __call__(self, x):
+        try:
+            return self.fn(x)
+        except Exception as e:
+            logger.error(f'[process_map] {type(e).__name__}: {e}')
+            return None
+
+
+def _auto_chunksize(n: int, workers: int | None) -> int:
+    '''``chunksize`` を自動で決める（``process_map`` で未指定のとき）。
+
+    子プロセスへは 1 件ずつより、まとめて送った方が速くなりやすい。そのまとめ数。
+
+    ``w`` は並列数。引数で決まっていなければ ``os.cpu_count()``、それも無ければ 4。
+    この **4** は「CPU が分からないときの仮の並列数」。式 ``n // (w * 4)`` の **4** とは別物。
+
+    ``n // (w * 4)`` の方の **4** は経験則の係数。ざっくり言うとチャンクの個数が
+    ``w * 4`` 前後になりやすく、負荷が均等ならワーカーあたりだいたい **4 回分の塊**
+    を処理するイメージ（厳密ではない）。
+
+    例: ``n=200``, ``w=5`` なら ``200 // 20 = 10`` が chunksize。全体は 20 チャンク、
+    5 人で割ると 1 人あたり平均 4 チャンク（各 10 件）。
+
+    結果は ``min(64, …)`` で上限。塊が大きすぎると **負荷が偏りやすい**。
+    タスクの重さがバラバラなとき、太い塊の中に遅いのが多く入ったワーカーだけが
+    長引き、他は先に終わって手待ちしがち（終盤のムラ）。塊を細かくすると配り直しの
+    機会が増えて和らぎやすい。進捗バーも細かく動きやすい。
+
+    ``max(1, …)`` で下限。割り算で 0 になっても最低 1 件は送る。
+    '''
+    w = workers or os.cpu_count() or 4
+    return max(1, min(64, n // (w * 4)))
+
+
+def process_map[T, R](
+    worker: Callable[[T], R],
+    items: Iterable[T],
+    workers: int | None = None,
+    *,
+    chunksize: int | None = None,
+) -> list[R | None]:
+    '''``ProcessPoolExecutor`` で ``worker`` を並列実行する。
+
+    子プロセスで例外が出た分は ``None`` で返す。全体は止めない。
+    進捗バーは常に tqdm。
+
+    ``chunksize`` は子へまとめて送る件数。省略なら自動。
+    進捗を細かくしたい・タスクの重さがバラバラで末尾に重いのが残る、なら ``chunksize=1``。
+    '''
+    safe = _SafeWorker(worker)
+    item_list = list(items)
+    cs = chunksize if chunksize is not None else _auto_chunksize(len(item_list), workers)
+    with ProcessPoolExecutor(max_workers=workers) as ex:
+        return list(
+            tqdm(ex.map(safe, item_list, chunksize=cs), total=len(item_list), unit='file')
+        )
+
+def glob_paths(dir_path: Path, pattern: str = '*.html') -> list[str]:
+    '''
+    ``dir_path`` 直下で ``pattern`` に一致するパスを ``str`` のリストで返す。
+
+    ``str`` にしているのは ``process_map`` 等のプロセスプールへ渡すとき pickle しやすくするため。
+    '''
+    return [str(p) for p in dir_path.glob(pattern)]
+
+
+def counter(start: int = 1) -> Iterator[int]:
+    n = start
+    while True:
+        yield n
+        n += 1