jerry-thomas 0.3.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

datapipeline/io/serializers.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, is_dataclass
+from typing import Any, Dict, Type
+
+from datapipeline.domain.sample import Sample
+
+
+class BaseSerializer:
+    payload_mode = "sample"
+
+    def serialize_payload(self, sample: Sample) -> Any:  # pragma: no cover - abstract
+        raise NotImplementedError
+
+
+class BaseJsonLineSerializer(BaseSerializer):
+    def __call__(self, sample: Sample) -> str:
+        data = self.serialize_payload(sample)
+        return json.dumps(data, ensure_ascii=False, default=str) + "\n"
+
+
+class SampleJsonLineSerializer(BaseJsonLineSerializer):
+    payload_mode = "sample"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_full_payload()
+
+
+class VectorJsonLineSerializer(BaseJsonLineSerializer):
+    payload_mode = "vector"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_vector_payload()
+
+
+class BasePrintSerializer(BaseSerializer):
+    def __call__(self, sample: Sample) -> str:
+        value = self.serialize_payload(sample)
+        return f"{value}\n"
+
+
+class SamplePrintSerializer(BasePrintSerializer):
+    payload_mode = "sample"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_full_payload()
+
+
+class VectorPrintSerializer(BasePrintSerializer):
+    payload_mode = "vector"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_vector_payload()
+
+
+class BaseCsvRowSerializer(BaseSerializer):
+    def __call__(self, sample: Sample) -> list[str | Any]:
+        key_value = sample.key
+        if isinstance(key_value, tuple):
+            key_struct = list(key_value)
+        else:
+            key_struct = key_value
+        if isinstance(key_struct, (list, dict)):
+            key_text = json.dumps(key_struct, ensure_ascii=False, default=str)
+        else:
+            key_text = "" if key_struct is None else str(key_struct)
+
+        payload_data = self.serialize_payload(sample)
+        if isinstance(payload_data, dict):
+            payload_data.pop("key", None)
+        payload_text = json.dumps(payload_data, ensure_ascii=False, default=str)
+        return [key_text, payload_text]
+
+
+class SampleCsvRowSerializer(BaseCsvRowSerializer):
+    payload_mode = "sample"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_full_payload()
+
+
+class VectorCsvRowSerializer(BaseCsvRowSerializer):
+    payload_mode = "vector"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_vector_payload()
+
+
+class BasePickleSerializer(BaseSerializer):
+    def __call__(self, sample: Sample) -> Any:
+        return self.serialize_payload(sample)
+
+
+class SamplePickleSerializer(BasePickleSerializer):
+    payload_mode = "sample"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample
+
+
+class VectorPickleSerializer(BasePickleSerializer):
+    payload_mode = "vector"
+
+    def serialize_payload(self, sample: Sample) -> Any:
+        return sample.as_vector_payload()
+
+
+def _record_payload(value: Any) -> Any:
+    if value is None:
+        return None
+    if is_dataclass(value):
+        return asdict(value)
+    if isinstance(value, dict):
+        return value
+    attrs = getattr(value, "__dict__", None)
+    if attrs:
+        return {
+            k: v
+            for k, v in attrs.items()
+            if not k.startswith("_")
+        }
+    return value
+
+
+def _record_key(value: Any) -> Any:
+    direct = getattr(value, "time", None)
+    if direct is not None:
+        return direct
+    record = getattr(value, "record", None)
+    if record is not None:
+        return getattr(record, "time", None)
+    return None
+
+
+class RecordJsonLineSerializer:
+    def __call__(self, record: Any) -> str:
+        payload = _record_payload(record)
+        return json.dumps(payload, ensure_ascii=False, default=str) + "\n"
+
+
+class RecordPrintSerializer:
+    def __call__(self, record: Any) -> str:
+        return f"{_record_payload(record)}\n"
+
+
+class RecordCsvRowSerializer:
+    def __call__(self, record: Any) -> list[str | Any]:
+        key_value = _record_key(record)
+        key_text = "" if key_value is None else str(key_value)
+        payload = json.dumps(_record_payload(record), ensure_ascii=False, default=str)
+        return [key_text, payload]
+
+
+class RecordPickleSerializer:
+    def __call__(self, record: Any) -> Any:
+        return record
+
+
+def _serializer_factory(
+    registry: Dict[str, Type[BaseSerializer]],
+    payload: str,
+    default_cls: Type[BaseSerializer],
+) -> BaseSerializer:
+    cls = registry.get(payload, default_cls)
+    return cls()
+
+
+JSON_SERIALIZERS: Dict[str, Type[BaseJsonLineSerializer]] = {
+    SampleJsonLineSerializer.payload_mode: SampleJsonLineSerializer,
+    VectorJsonLineSerializer.payload_mode: VectorJsonLineSerializer,
+}
+
+PRINT_SERIALIZERS: Dict[str, Type[BasePrintSerializer]] = {
+    SamplePrintSerializer.payload_mode: SamplePrintSerializer,
+    VectorPrintSerializer.payload_mode: VectorPrintSerializer,
+}
+
+CSV_SERIALIZERS: Dict[str, Type[BaseCsvRowSerializer]] = {
+    SampleCsvRowSerializer.payload_mode: SampleCsvRowSerializer,
+    VectorCsvRowSerializer.payload_mode: VectorCsvRowSerializer,
+}
+
+PICKLE_SERIALIZERS: Dict[str, Type[BasePickleSerializer]] = {
+    SamplePickleSerializer.payload_mode: SamplePickleSerializer,
+    VectorPickleSerializer.payload_mode: VectorPickleSerializer,
+}
+
+
+def json_line_serializer(payload: str) -> BaseJsonLineSerializer:
+    return _serializer_factory(JSON_SERIALIZERS, payload, SampleJsonLineSerializer)
+
+
+def print_serializer(payload: str) -> BasePrintSerializer:
+    return _serializer_factory(PRINT_SERIALIZERS, payload, SamplePrintSerializer)
+
+
+def csv_row_serializer(payload: str) -> BaseCsvRowSerializer:
+    return _serializer_factory(CSV_SERIALIZERS, payload, SampleCsvRowSerializer)
+
+
+def pickle_serializer(payload: str) -> BasePickleSerializer:
+    return _serializer_factory(PICKLE_SERIALIZERS, payload, SamplePickleSerializer)
+
+
+def record_json_line_serializer() -> RecordJsonLineSerializer:
+    return RecordJsonLineSerializer()
+
+
+def record_print_serializer() -> RecordPrintSerializer:
+    return RecordPrintSerializer()
+
+
+def record_csv_row_serializer() -> RecordCsvRowSerializer:
+    return RecordCsvRowSerializer()
+
+
+def record_pickle_serializer() -> RecordPickleSerializer:
+    return RecordPickleSerializer()

datapipeline/io/sinks/__init__.py

@@ -0,0 +1,23 @@
+from .base import BaseSink
+from .stdout import StdoutTextSink
+from .rich import (
+    RichFormatter,
+    ReprRichFormatter,
+    JsonRichFormatter,
+    PlainRichFormatter,
+    RichStdoutSink,
+)
+from .files import AtomicTextFileSink, AtomicBinaryFileSink, GzipBinarySink
+
+__all__ = [
+    "BaseSink",
+    "StdoutTextSink",
+    "RichFormatter",
+    "ReprRichFormatter",
+    "JsonRichFormatter",
+    "PlainRichFormatter",
+    "RichStdoutSink",
+    "AtomicTextFileSink",
+    "AtomicBinaryFileSink",
+    "GzipBinarySink",
+]

datapipeline/io/sinks/base.py

@@ -0,0 +1,2 @@
+class BaseSink:
+    def close(self) -> None: ...

datapipeline/io/sinks/files.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+import os
+import tempfile
+import gzip
+
+from .base import BaseSink
+
+
+class AtomicTextFileSink(BaseSink):
+    def __init__(self, dest: Path):
+        self._dest = dest
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        self._tmp = Path(
+            tempfile.NamedTemporaryFile(dir=str(dest.parent), delete=False).name
+        )
+        self._fh = open(self._tmp, "w", encoding="utf-8")
+
+    @property
+    def file_path(self) -> Path:
+        return self._dest
+
+    @property
+    def fh(self):
+        return self._fh
+
+    def write_text(self, s: str) -> None:
+        self._fh.write(s)
+
+    def close(self) -> None:
+        self._fh.close()
+        os.replace(self._tmp, self._dest)
+
+
+class AtomicBinaryFileSink(BaseSink):
+    def __init__(self, dest: Path):
+        self._dest = dest
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        self._tmp = Path(
+            tempfile.NamedTemporaryFile(dir=str(dest.parent), delete=False).name
+        )
+        self._fh = open(self._tmp, "wb")
+
+    @property
+    def file_path(self) -> Path:
+        return self._dest
+
+    @property
+    def fh(self):
+        return self._fh
+
+    def write_bytes(self, b: bytes) -> None:
+        self._fh.write(b)
+
+    def close(self) -> None:
+        self._fh.close()
+        os.replace(self._tmp, self._dest)
+
+
+class GzipBinarySink(BaseSink):
+    def __init__(self, dest: Path):
+        self._dest = dest
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        self._tmp = Path(
+            tempfile.NamedTemporaryFile(dir=str(dest.parent), delete=False).name
+        )
+        self._raw = open(self._tmp, "wb")
+        self._fh = gzip.GzipFile(fileobj=self._raw, mode="wb")
+
+    @property
+    def file_path(self) -> Path:
+        return self._dest
+
+    def write_bytes(self, b: bytes) -> None:
+        self._fh.write(b)
+
+    def close(self) -> None:
+        self._fh.close()
+        self._raw.close()
+        os.replace(self._tmp, self._dest)

datapipeline/io/sinks/rich.py

@@ -0,0 +1,57 @@
+from typing import Protocol
+
+from .stdout import StdoutTextSink
+
+
+class RichFormatter(Protocol):
+    def render(self, console, text: str) -> None: ...
+
+
+class ReprRichFormatter:
+    def __init__(self):
+        from rich.highlighter import ReprHighlighter
+
+        self._highlighter = ReprHighlighter()
+
+    def render(self, console, text: str) -> None:
+        console.print(self._highlighter(text))
+
+
+class JsonRichFormatter:
+    def render(self, console, text: str) -> None:
+        import json as _json
+
+        stripped = text.strip()
+        if not stripped:
+            return
+        try:
+            data = _json.loads(stripped)
+            console.print_json(data=data)
+        except Exception:
+            console.print(stripped)
+
+
+class PlainRichFormatter:
+    def render(self, console, text: str) -> None:
+        console.print(text)
+
+
+class RichStdoutSink(StdoutTextSink):
+    def __init__(self, formatter: RichFormatter):
+        super().__init__()
+        try:
+            from rich.console import Console
+        except Exception:  # pragma: no cover
+            self.console = None
+        else:
+            self.console = Console(
+                file=self.stream, markup=False, highlight=False, soft_wrap=True
+            )
+        self._formatter = formatter
+
+    def write_text(self, s: str) -> None:
+        if not self.console:
+            super().write_text(s)
+            return
+        text = s.rstrip("\n")
+        self._formatter.render(self.console, text)

datapipeline/io/sinks/stdout.py

@@ -0,0 +1,18 @@
+import sys
+from typing import Optional
+
+from .base import BaseSink
+
+
+class StdoutTextSink(BaseSink):
+    def __init__(self, stream: Optional[object] = None):
+        self.stream = stream or sys.stdout
+
+    def write_text(self, s: str) -> None:
+        self.stream.write(s)
+
+    def flush(self) -> None:
+        self.stream.flush()
+
+    def close(self) -> None:
+        self.flush()

datapipeline/io/writers/__init__.py

@@ -0,0 +1,14 @@
+from .base import LineWriter, HeaderJsonlMixin
+from .jsonl import JsonLinesStdoutWriter, JsonLinesFileWriter, GzipJsonLinesWriter
+from .csv_writer import CsvFileWriter
+from .pickle_writer import PickleFileWriter
+
+__all__ = [
+    "LineWriter",
+    "HeaderJsonlMixin",
+    "JsonLinesStdoutWriter",
+    "JsonLinesFileWriter",
+    "GzipJsonLinesWriter",
+    "CsvFileWriter",
+    "PickleFileWriter",
+]

datapipeline/io/writers/base.py

@@ -0,0 +1,28 @@
+import json
+from typing import Optional
+
+from datapipeline.io.protocols import HeaderCapable, Writer
+from datapipeline.io.sinks import StdoutTextSink, AtomicTextFileSink
+
+
+class LineWriter(Writer):
+    """Text line writer (uses a text sink + serializer)."""
+
+    def __init__(self, sink: StdoutTextSink | AtomicTextFileSink, serializer):
+        self.sink = sink
+        self.serializer = serializer
+
+    def write(self, item) -> None:
+        self.sink.write_text(self.serializer(item))
+
+    def close(self) -> None:
+        self.sink.close()
+
+
+class HeaderJsonlMixin(HeaderCapable):
+    """Provide a header write by emitting one JSON line."""
+
+    def write_header(self, header: dict) -> None:
+        self.sink.write_text(
+            json.dumps({"__checkpoint__": header}, ensure_ascii=False) + "\n"
+        )

datapipeline/io/writers/csv_writer.py

@@ -0,0 +1,25 @@
+import csv
+from pathlib import Path
+from typing import Optional
+
+from datapipeline.io.serializers import csv_row_serializer, BaseCsvRowSerializer
+from datapipeline.io.protocols import HasFilePath, Writer
+from datapipeline.io.sinks import AtomicTextFileSink
+
+
+class CsvFileWriter(Writer, HasFilePath):
+    def __init__(self, dest: Path, serializer: BaseCsvRowSerializer | None = None):
+        self.sink = AtomicTextFileSink(dest)
+        self.writer = csv.writer(self.sink.fh)
+        self.writer.writerow(["key", "values"])
+        self._serializer = serializer or csv_row_serializer("sample")
+
+    @property
+    def file_path(self) -> Optional[Path]:
+        return self.sink.file_path
+
+    def write(self, item) -> None:
+        self.writer.writerow(self._serializer(item))
+
+    def close(self) -> None:
+        self.sink.close()

datapipeline/io/writers/jsonl.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import Optional
+import json
+
+from datapipeline.io.serializers import json_line_serializer, BaseJsonLineSerializer
+from datapipeline.io.protocols import HeaderCapable, HasFilePath, Writer
+from datapipeline.io.sinks import (
+    AtomicTextFileSink,
+    GzipBinarySink,
+    StdoutTextSink,
+)
+
+from .base import HeaderJsonlMixin, LineWriter
+
+
+class JsonLinesStdoutWriter(LineWriter, HeaderJsonlMixin):
+    def __init__(self, serializer: BaseJsonLineSerializer | None = None):
+        super().__init__(StdoutTextSink(), serializer or json_line_serializer("sample"))
+
+
+class JsonLinesFileWriter(LineWriter, HeaderJsonlMixin, HasFilePath):
+    def __init__(self, dest: Path, serializer: BaseJsonLineSerializer | None = None):
+        self._sink = AtomicTextFileSink(dest)
+        super().__init__(self._sink, serializer or json_line_serializer("sample"))
+
+    @property
+    def file_path(self) -> Optional[Path]:
+        return self._sink.file_path
+
+
+class GzipJsonLinesWriter(Writer, HeaderCapable, HasFilePath):
+    def __init__(self, dest: Path, serializer: BaseJsonLineSerializer | None = None):
+        self.sink = GzipBinarySink(dest)
+        self._serializer = serializer or json_line_serializer("sample")
+
+    @property
+    def file_path(self) -> Optional[Path]:
+        return self.sink.file_path
+
+    def write_header(self, header: dict) -> None:
+        self.sink.write_bytes(
+            (json.dumps({"__checkpoint__": header}, ensure_ascii=False) + "\n").encode(
+                "utf-8"
+            )
+        )
+
+    def write(self, item) -> None:
+        line = self._serializer(item)
+        self.sink.write_bytes(line.encode("utf-8"))
+
+    def close(self) -> None:
+        self.sink.close()

datapipeline/io/writers/pickle_writer.py

@@ -0,0 +1,30 @@
+import pickle
+from pathlib import Path
+from typing import Optional
+
+from datapipeline.io.serializers import pickle_serializer, BasePickleSerializer
+from datapipeline.io.protocols import HasFilePath, Writer
+from datapipeline.io.sinks import AtomicBinaryFileSink
+
+
+class PickleFileWriter(Writer, HasFilePath):
+    def __init__(
+        self,
+        dest: Path,
+        serializer: BasePickleSerializer | None = None,
+        protocol: int = pickle.HIGHEST_PROTOCOL,
+    ):
+        self.sink = AtomicBinaryFileSink(dest)
+        self.pickler = pickle.Pickler(self.sink.fh, protocol=protocol)
+        self._serializer = serializer or pickle_serializer("sample")
+
+    @property
+    def file_path(self) -> Optional[Path]:
+        return self.sink.file_path
+
+    def write(self, item) -> None:
+        self.pickler.dump(self._serializer(item))
+        self.pickler.clear_memo()
+
+    def close(self) -> None:
+        self.sink.close()

datapipeline/pipeline/artifacts.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+from datapipeline.config.dataset.dataset import FeatureDatasetConfig
+from datapipeline.config.dataset.feature import FeatureRecordConfig
+from datapipeline.services.constants import (
+    SCALER_STATISTICS,
+    VECTOR_SCHEMA,
+    VECTOR_SCHEMA_METADATA,
+)
+
+
+@dataclass(frozen=True)
+class StageDemand:
+    stage: int | None
+
+
+def _needs_scaler(configs: Iterable[FeatureRecordConfig]) -> bool:
+    for cfg in configs:
+        scale = getattr(cfg, "scale", False)
+        if isinstance(scale, dict):
+            return True
+        if bool(scale):
+            return True
+    return False
+
+
+def _requires_scaler(dataset: FeatureDatasetConfig) -> bool:
+    if _needs_scaler(dataset.features or []):
+        return True
+    if dataset.targets:
+        return _needs_scaler(dataset.targets)
+    return False
+
+
+def required_artifacts_for(
+    dataset: FeatureDatasetConfig,
+    demands: Iterable[StageDemand],
+) -> set[str]:
+    required: set[str] = set()
+    needs_metadata = False
+    for demand in demands:
+        stage = demand.stage
+        effective_stage = 7 if stage is None else stage
+
+        if effective_stage >= 5 and _requires_scaler(dataset):
+            required.add(SCALER_STATISTICS)
+
+        if effective_stage >= 6:
+            required.add(VECTOR_SCHEMA)
+            needs_metadata = True
+
+    if needs_metadata:
+        required.add(VECTOR_SCHEMA_METADATA)
+
+    return required