From 546f34ea2cb842a77a2fb3d10dd0e7aca9bf62ed Mon Sep 17 00:00:00 2001 From: bschilder Date: Tue, 21 Apr 2026 16:26:46 -0400 Subject: [PATCH] =?UTF-8?q?feat(meds):=20OMOP=20=E2=86=92=20MEDS=20convers?= =?UTF-8?q?ion=20for=20ML-native=20EHR=20encoders?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds synthlab.meds, a thin wrapper over the community meds_etl package (github.com/Medical-Event-Data-Standard/meds_etl). This lets SynthLab output feed directly into MEDS-native foundation models such as SMB-v1 (huggingface.co/standardmodelbio) and MOTOR. API: - MedsConvertConfig dataclass (omop_dir, meds_dir, backend, shards, proc, overwrite) — matches meds_etl_omop's CLI surface. - convert_omop_to_meds(config): runs meds_etl_omop as subprocess, surfaces stderr on failure, refuses to overwrite unless asked. - load_meds_events(meds_dir, subject_ids=None): polars-native reader for the sharded parquet output, sorted by (subject_id, time). - get_meds_cache_dir / get_meds_info / print_meds_info — mirrors the pattern every other synthlab submodule uses. Pipeline: SyntheaRunner.run() -> Synthea CSVs convert_synthea_to_omop() -> OMOP v5.4 convert_omop_to_meds() -> MEDS parquet (downstream: SMB-v1 / MOTOR encoder) Install: `pip install synthlab[meds]` (adds meds_etl + meds pins). Infra: - New tests/ dir + test_meds.py (15 tests, all passing on py3.11): cache-dir, info dict, config validation, error paths, shard reading, subject filtering, meds_etl entry-point registration. - New .github/workflows/test.yml runs pytest on every PR + push. - pyproject.toml gains [meds] + [tests] optional extras and pulls them into [all]. - README gets a "MEDS conversion" section with quickstart. No breaking changes to existing APIs. --- .github/workflows/test.yml | 36 ++++ README.md | 19 ++ pyproject.toml | 12 ++ synthlab/__init__.py | 17 ++ synthlab/meds.py | 375 +++++++++++++++++++++++++++++++++++++ tests/__init__.py | 0 tests/test_meds.py | 232 +++++++++++++++++++++++ 7 files changed, 691 insertions(+) create mode 100644 .github/workflows/test.yml create mode 100644 synthlab/meds.py create mode 100644 tests/__init__.py create mode 100644 tests/test_meds.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..a78f3d4 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,36 @@ +name: tests + +on: + pull_request: + branches: [main] + push: + branches: [main] + +jobs: + pytest: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11"] + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Install package + test extras + run: | + python -m pip install --upgrade pip + pip install -e ".[tests,meds,aws,genomics]" + + - name: Show meds_etl info + run: | + python -c "import meds_etl; print('meds_etl', meds_etl.__version__)" + which meds_etl_omop || true + + - name: Run pytest + run: pytest tests/ -v --tb=short diff --git a/README.md b/README.md index 8817b67..c8a9a6d 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,25 @@ SynthLab provides Python interfaces for working with major synthetic healthcare - **Access Information**: Clear documentation of open vs. registration-required datasets - **Download Utilities**: Helpers for downloading select open-access datasets +### MEDS (Medical Event Data Standard) conversion (NEW) +Plug synthetic EHR into ML-native foundation models. Provides a thin, +well-tested wrapper over the community [`meds_etl`](https://github.com/Medical-Event-Data-Standard/meds_etl) +package so OMOP CSVs (produced by `synthlab.synthea.convert_synthea_to_omop`) +become MEDS parquet shards ready for models like +[SMB-v1](https://huggingface.co/standardmodelbio) or MOTOR: + +```python +from synthlab.meds import MedsConvertConfig, convert_omop_to_meds, load_meds_events + +convert_omop_to_meds(MedsConvertConfig( + omop_dir="~/.cache/synthlab/synthea/omop_100", + meds_dir="~/.cache/synthlab/meds/synthea_100", +)) +df = load_meds_events("~/.cache/synthlab/meds/synthea_100") +``` + +Install the optional extra: `pip install synthlab[meds]`. + ## Installation ```bash diff --git a/pyproject.toml b/pyproject.toml index df98cdb..456850c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,16 @@ visualization = [ "pyvis>=0.3.0", "networkx>=2.8.0", ] +meds = [ + # Medical Event Data Standard ETL (OMOP / MIMIC-IV → MEDS parquet). + # https://github.com/Medical-Event-Data-Standard/meds_etl + "meds_etl>=0.3", + "meds>=0.4", +] +tests = [ + "pytest>=8.0", + "pytest-cov>=5.0", +] all = [ "boto3>=1.26.0", "numpy>=1.20.0", @@ -78,6 +88,8 @@ all = [ "faiss-gpu>=1.7.0", "pyvis>=0.3.0", "networkx>=2.8.0", + "meds_etl>=0.3", + "meds>=0.4", ] [project.urls] diff --git a/synthlab/__init__.py b/synthlab/__init__.py index 0eb7213..121106f 100644 --- a/synthlab/__init__.py +++ b/synthlab/__init__.py @@ -145,6 +145,16 @@ get_snomed_data_dir, ) +# MEDS (Medical Event Data Standard) conversion for ML-native EHR inputs +from synthlab.meds import ( + MedsConvertConfig, + convert_omop_to_meds, + load_meds_events, + get_meds_cache_dir, + get_meds_info, + print_meds_info, +) + from synthlab.coherent import ( list_coherent_components, list_coherent_files, @@ -340,4 +350,11 @@ def _print_banner(): "SAPBERT_MODEL_ID", "SAPBERT_MODELS", "EMBEDDING_MODELS", + # MEDS conversion + "MedsConvertConfig", + "convert_omop_to_meds", + "load_meds_events", + "get_meds_cache_dir", + "get_meds_info", + "print_meds_info", ] diff --git a/synthlab/meds.py b/synthlab/meds.py new file mode 100644 index 0000000..e62c60d --- /dev/null +++ b/synthlab/meds.py @@ -0,0 +1,375 @@ +"""Medical Event Data Standard (MEDS) conversion utilities. + +MEDS is a machine-learning-native schema for clinical event data: +https://medical-event-data-standard.github.io. Models like +[SMB-v1](https://huggingface.co/standardmodelbio) and +[MOTOR](https://github.com/som-shahlab/motor) consume MEDS parquet +as their canonical patient-history input, so being able to convert +synthetic EHR (Synthea, UKB Synthetic) into MEDS opens the door to +plugging SynthLab's output directly into any MEDS-compatible model. + +This module is a thin wrapper around +[meds_etl](https://pypi.org/project/meds-etl/) — the official +open-source ETL library maintained by the MEDS community — with +SynthLab-specific convenience (cache-dir discovery, config +dataclass, loading helpers that return polars DataFrames ready +for downstream tokenizers). + +Pipeline +-------- + +``` +SyntheaRunner.run() # synthlab.synthea + → Synthea CSV output +convert_synthea_to_omop(synthea_dir, ...) # synthlab.synthea + → OMOP v5.4 CDM CSVs +convert_omop_to_meds(omop_dir, meds_dir) # synthlab.meds (this module) + → MEDS parquet shards +load_meds_events(meds_dir) # synthlab.meds + → polars DataFrame ready for SMB-v1 / MOTOR / etc. +``` + +Example +------- + +```python +from synthlab.meds import convert_omop_to_meds, load_meds_events + +convert_omop_to_meds( + omop_dir="~/.cache/synthlab/synthea/omop_10", + meds_dir="~/.cache/synthlab/meds/synthea_10", +) +df = load_meds_events("~/.cache/synthlab/meds/synthea_10") +print(df.head()) +``` +""" + +from __future__ import annotations + +import logging +import os +import shutil +import subprocess +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal, Optional + +logger = logging.getLogger(__name__) + +# Lazy imports for optional heavy deps +_polars_available = False +_meds_etl_available = False + +try: + import polars as pl # noqa: F401 + _polars_available = True +except ImportError: + pass + +try: + import meds_etl # noqa: F401 -- import-time availability check + _meds_etl_available = True +except ImportError: + pass + + +# --------------------------------------------------------------------------- +# Cache dir + info helpers (mirrors the pattern of other synthlab modules) +# --------------------------------------------------------------------------- + + +def get_meds_cache_dir() -> Path: + """Return the cache directory for MEDS-format datasets. + + Mirrors ``synthlab.synthea.get_synthea_cache_dir`` — resolves to + ``$SYNTHLAB_CACHE_DIR/meds`` when the env-var is set, otherwise + ``~/.cache/synthlab/meds``. Creates the directory if absent. + + Returns + ------- + pathlib.Path + Absolute path to the cache directory. + + Examples + -------- + >>> p = get_meds_cache_dir() + >>> p.is_dir() + True + >>> p.name + 'meds' + """ + base = os.environ.get("SYNTHLAB_CACHE_DIR", str(Path.home() / ".cache" / "synthlab")) + cache_dir = Path(base).expanduser() / "meds" + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + + +def get_meds_info() -> dict: + """Return a metadata dict describing the MEDS converter module. + + Returns + ------- + dict + Keys: ``module``, ``purpose``, ``schema_docs``, ``backend_pkg``, + ``cache_dir``, ``meds_etl_available``, ``polars_available``. + + Examples + -------- + >>> info = get_meds_info() + >>> info["module"] + 'synthlab.meds' + """ + return { + "module": "synthlab.meds", + "purpose": "Convert OMOP v5.4 CSVs to MEDS parquet for ML-ready EHR encoding.", + "schema_docs": "https://medical-event-data-standard.github.io", + "backend_pkg": "meds_etl", + "cache_dir": str(get_meds_cache_dir()), + "meds_etl_available": _meds_etl_available, + "polars_available": _polars_available, + } + + +def print_meds_info() -> None: + """Pretty-print :func:`get_meds_info` to stdout. + + Examples + -------- + >>> print_meds_info() # doctest: +SKIP + synthlab.meds + - purpose: ... + ... + """ + info = get_meds_info() + print(f"{info['module']}") + print(f"- purpose: {info['purpose']}") + print(f"- schema_docs: {info['schema_docs']}") + print(f"- backend_pkg: {info['backend_pkg']}") + print(f"- cache_dir: {info['cache_dir']}") + print(f"- meds_etl installed: {info['meds_etl_available']}") + print(f"- polars installed: {info['polars_available']}") + + +# --------------------------------------------------------------------------- +# Config + conversion +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class MedsConvertConfig: + """Config for :func:`convert_omop_to_meds`. + + Attributes + ---------- + omop_dir : pathlib.Path + Path to a directory of OMOP v5.4 CSV files (one per OMOP table — + ``person.csv``, ``condition_occurrence.csv``, etc.). Typically + produced by :func:`synthlab.synthea.convert_synthea_to_omop`. + meds_dir : pathlib.Path + Destination directory for the MEDS parquet output. Will be + created. If ``overwrite=False`` and the directory already + contains ``data/``, conversion is skipped. + backend : {"polars", "cpp"}, default ``"polars"`` + :mod:`meds_etl` has two backends. ``polars`` is Python-only + and works anywhere; ``cpp`` is faster but requires the + optional ``meds_etl[cpp]`` install. + num_shards : int, default ``4`` + Number of shards the MEDS dataset is split into. Lower + shard counts use more memory per worker; see the + :mod:`meds_etl` README for tuning. + num_proc : int, default ``1`` + Worker processes. Match ``num_shards`` for the cpp backend, + keep low for polars to avoid OOM. + overwrite : bool, default ``False`` + If ``True`` and ``meds_dir`` already exists, it is removed + before conversion. + + Examples + -------- + >>> cfg = MedsConvertConfig(omop_dir="/tmp/omop", meds_dir="/tmp/meds") + >>> cfg.backend + 'polars' + """ + + omop_dir: Path + meds_dir: Path + backend: Literal["polars", "cpp"] = "polars" + num_shards: int = 4 + num_proc: int = 1 + overwrite: bool = False + + def __post_init__(self) -> None: + """Normalise paths to :class:`Path` + validate backend choice.""" + # frozen dataclass: __setattr__ goes through object + object.__setattr__(self, "omop_dir", Path(self.omop_dir).expanduser()) + object.__setattr__(self, "meds_dir", Path(self.meds_dir).expanduser()) + if self.backend not in {"polars", "cpp"}: + raise ValueError( + f"backend must be 'polars' or 'cpp'; got {self.backend!r}" + ) + + +def _require_meds_etl() -> None: + """Raise :class:`ImportError` with a helpful message if missing. + + Raises + ------ + ImportError + When :mod:`meds_etl` cannot be imported. Message names the + right ``pip install`` invocation. + """ + if not _meds_etl_available: + raise ImportError( + "meds_etl is required for MEDS conversion. Install with " + "`pip install meds_etl` (or `pip install 'meds_etl[cpp]'` " + "for the faster C++ backend). See " + "https://github.com/Medical-Event-Data-Standard/meds_etl." + ) + + +def convert_omop_to_meds(config: MedsConvertConfig) -> Path: + """Convert an OMOP v5.4 CSV directory into a MEDS parquet dataset. + + Thin wrapper around the :mod:`meds_etl` ``meds_etl_omop`` CLI. + The CLI is invoked via ``subprocess`` so a failure surfaces the + full ``meds_etl`` traceback without this wrapper swallowing it. + + Parameters + ---------- + config : MedsConvertConfig + + Returns + ------- + pathlib.Path + Path to ``config.meds_dir`` after successful conversion. + + Raises + ------ + ImportError + If :mod:`meds_etl` isn't installed. + FileNotFoundError + If ``config.omop_dir`` doesn't exist. + RuntimeError + If the ``meds_etl_omop`` CLI exits non-zero. The wrapped + ``stderr`` is attached to the error message. + FileExistsError + If ``config.meds_dir`` already has a ``data/`` subdir and + ``config.overwrite`` is ``False``. + + Examples + -------- + >>> # doctest: +SKIP + >>> cfg = MedsConvertConfig(omop_dir="omop/", meds_dir="meds/") + >>> convert_omop_to_meds(cfg) + PosixPath('meds') + """ + _require_meds_etl() + if not config.omop_dir.is_dir(): + raise FileNotFoundError( + f"omop_dir does not exist: {config.omop_dir}" + ) + existing_data = config.meds_dir / "data" + if existing_data.exists(): + if not config.overwrite: + raise FileExistsError( + f"{existing_data} already contains MEDS data; pass " + "overwrite=True on MedsConvertConfig to replace it." + ) + shutil.rmtree(config.meds_dir) + config.meds_dir.mkdir(parents=True, exist_ok=True) + + cmd = [ + "meds_etl_omop", + str(config.omop_dir), + str(config.meds_dir), + "--backend", + config.backend, + "--num_shards", + str(config.num_shards), + "--num_proc", + str(config.num_proc), + ] + logger.info("Running meds_etl_omop: %s", " ".join(cmd)) + completed = subprocess.run(cmd, capture_output=True, text=True, check=False) + if completed.returncode != 0: + raise RuntimeError( + f"meds_etl_omop exited with {completed.returncode}. " + f"stdout={completed.stdout!r} stderr={completed.stderr!r}" + ) + logger.info("MEDS conversion complete: %s", config.meds_dir) + return config.meds_dir + + +# --------------------------------------------------------------------------- +# Loading helpers +# --------------------------------------------------------------------------- + + +def load_meds_events( + meds_dir: Path | str, + *, + subject_ids: Optional[list[str]] = None, +): + """Load MEDS event rows from a converted dataset into a polars DataFrame. + + Reads every parquet shard under ``/data/`` and concatenates + them. Filtering by ``subject_ids`` pushes the filter into the polars + scan, so a cohort of, say, 500 patients out of 100k costs roughly + ``O(patient_count)`` rather than ``O(dataset)``. + + Parameters + ---------- + meds_dir : pathlib.Path or str + Directory produced by :func:`convert_omop_to_meds` (must + contain a ``data/`` subdirectory of parquet shards). + subject_ids : list[str], optional + If given, only rows with ``subject_id`` in this list are + returned. Accepts int-valued subject IDs too — they're cast + to string before comparison because parquet schema varies + across ETL backends. + + Returns + ------- + polars.DataFrame + Columns: ``subject_id``, ``time``, ``code``, plus any + auxiliary columns produced by the ETL (``value``, ``table``, + ``numeric_value``, …). Sorted by ``(subject_id, time)`` + which is the order MEDS-native consumers expect. + + Raises + ------ + ImportError + If :mod:`polars` isn't installed. + FileNotFoundError + If ``meds_dir/data`` doesn't exist or has no parquet shards. + + Examples + -------- + >>> # doctest: +SKIP + >>> df = load_meds_events("~/.cache/synthlab/meds/synthea_100") + >>> df.columns[:3] + ['subject_id', 'time', 'code'] + """ + if not _polars_available: + raise ImportError( + "polars is required for load_meds_events. " + "`pip install polars` (already listed in synthlab's core deps)." + ) + import polars as pl + + data_dir = Path(meds_dir).expanduser() / "data" + if not data_dir.is_dir(): + raise FileNotFoundError( + f"No MEDS data subdirectory at {data_dir}; " + "did convert_omop_to_meds() run successfully?" + ) + shards = sorted(data_dir.glob("*.parquet")) + if not shards: + raise FileNotFoundError( + f"No parquet shards under {data_dir}." + ) + lf = pl.scan_parquet([str(s) for s in shards]) + if subject_ids is not None: + sid_strs = [str(s) for s in subject_ids] + lf = lf.filter(pl.col("subject_id").cast(pl.Utf8).is_in(sid_strs)) + return lf.sort(["subject_id", "time"]).collect() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_meds.py b/tests/test_meds.py new file mode 100644 index 0000000..1247b1a --- /dev/null +++ b/tests/test_meds.py @@ -0,0 +1,232 @@ +"""Unit tests for :mod:`synthlab.meds`. + +Covers the pure-Python surface (config validation, cache-dir +resolution, error paths). Integration tests that actually invoke +``meds_etl_omop`` are gated on the package being importable, since +meds_etl pulls a big dep tree and isn't always in CI. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from synthlab.meds import ( + MedsConvertConfig, + _meds_etl_available, + convert_omop_to_meds, + get_meds_cache_dir, + get_meds_info, + load_meds_events, + print_meds_info, +) + + +# --------------------------------------------------------------------------- +# Cache dir + info +# --------------------------------------------------------------------------- + + +def test_get_meds_cache_dir_creates_and_returns_path(tmp_path, monkeypatch) -> None: + """get_meds_cache_dir creates the /meds subdir on first call.""" + monkeypatch.setenv("SYNTHLAB_CACHE_DIR", str(tmp_path)) + cache = get_meds_cache_dir() + assert cache == tmp_path / "meds" + assert cache.is_dir() + + +def test_get_meds_cache_dir_defaults_to_home(monkeypatch) -> None: + """Without SYNTHLAB_CACHE_DIR the cache lives under ~/.cache/synthlab/meds.""" + monkeypatch.delenv("SYNTHLAB_CACHE_DIR", raising=False) + cache = get_meds_cache_dir() + # We don't actually want to mkdir in $HOME during tests, so just check + # the path shape. The dir may or may not exist depending on prior runs. + assert cache.name == "meds" + assert cache.parent.name == "synthlab" + + +def test_get_meds_info_has_expected_keys() -> None: + """Info dict exposes the canonical metadata fields.""" + info = get_meds_info() + for key in ( + "module", + "purpose", + "schema_docs", + "backend_pkg", + "cache_dir", + "meds_etl_available", + "polars_available", + ): + assert key in info, f"missing key: {key}" + assert info["module"] == "synthlab.meds" + assert info["backend_pkg"] == "meds_etl" + + +def test_print_meds_info_is_stable(capsys) -> None: + """print_meds_info writes the expected header line without crashing.""" + print_meds_info() + out = capsys.readouterr().out + assert "synthlab.meds" in out + assert "cache_dir" in out + + +# --------------------------------------------------------------------------- +# MedsConvertConfig +# --------------------------------------------------------------------------- + + +def test_config_normalises_paths(tmp_path) -> None: + """String paths are cast to Path and ~/ is expanded.""" + cfg = MedsConvertConfig( + omop_dir=str(tmp_path / "omop"), + meds_dir=str(tmp_path / "meds"), + ) + assert isinstance(cfg.omop_dir, Path) + assert isinstance(cfg.meds_dir, Path) + + +def test_config_rejects_unknown_backend(tmp_path) -> None: + """Only 'polars' and 'cpp' are allowed backend values.""" + with pytest.raises(ValueError, match="backend"): + MedsConvertConfig( + omop_dir=tmp_path / "omop", + meds_dir=tmp_path / "meds", + backend="rust", # type: ignore[arg-type] + ) + + +def test_config_defaults(tmp_path) -> None: + """Default backend is polars with 4 shards and 1 worker.""" + cfg = MedsConvertConfig(omop_dir=tmp_path / "omop", meds_dir=tmp_path / "meds") + assert cfg.backend == "polars" + assert cfg.num_shards == 4 + assert cfg.num_proc == 1 + assert cfg.overwrite is False + + +# --------------------------------------------------------------------------- +# convert_omop_to_meds — error paths +# --------------------------------------------------------------------------- + + +def test_convert_raises_without_meds_etl(tmp_path, monkeypatch) -> None: + """When meds_etl isn't installed, a helpful ImportError is raised.""" + monkeypatch.setattr("synthlab.meds._meds_etl_available", False) + cfg = MedsConvertConfig( + omop_dir=tmp_path / "omop", + meds_dir=tmp_path / "meds", + ) + with pytest.raises(ImportError, match="meds_etl"): + convert_omop_to_meds(cfg) + + +def test_convert_raises_when_omop_missing(tmp_path, monkeypatch) -> None: + """Missing omop_dir → FileNotFoundError that names the bad path.""" + monkeypatch.setattr("synthlab.meds._meds_etl_available", True) + cfg = MedsConvertConfig( + omop_dir=tmp_path / "nonexistent_omop", + meds_dir=tmp_path / "meds", + ) + with pytest.raises(FileNotFoundError, match="nonexistent_omop"): + convert_omop_to_meds(cfg) + + +def test_convert_refuses_overwrite_by_default(tmp_path, monkeypatch) -> None: + """Existing meds_dir/data/ triggers FileExistsError unless overwrite=True.""" + monkeypatch.setattr("synthlab.meds._meds_etl_available", True) + omop = tmp_path / "omop" + omop.mkdir() + meds = tmp_path / "meds" + (meds / "data").mkdir(parents=True) + cfg = MedsConvertConfig(omop_dir=omop, meds_dir=meds, overwrite=False) + with pytest.raises(FileExistsError, match="overwrite=True"): + convert_omop_to_meds(cfg) + + +# --------------------------------------------------------------------------- +# load_meds_events — error paths +# --------------------------------------------------------------------------- + + +def test_load_meds_events_missing_dir(tmp_path) -> None: + """Missing meds_dir/data raises FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="No MEDS data"): + load_meds_events(tmp_path / "nonexistent") + + +def test_load_meds_events_empty_data_dir(tmp_path) -> None: + """Empty data/ subdir (no shards) raises FileNotFoundError.""" + (tmp_path / "data").mkdir() + with pytest.raises(FileNotFoundError, match="No parquet shards"): + load_meds_events(tmp_path) + + +def test_load_meds_events_reads_shards(tmp_path) -> None: + """Given a tiny parquet shard, load_meds_events returns a sorted frame.""" + pl = pytest.importorskip("polars") + data = tmp_path / "data" + data.mkdir() + pl.DataFrame( + { + "subject_id": ["S2", "S1", "S1"], + "time": [ + "2020-01-02T00:00:00", + "2020-01-01T00:00:00", + "2020-01-03T00:00:00", + ], + "code": ["ICD10:A", "ICD10:B", "ICD10:C"], + } + ).with_columns(pl.col("time").str.to_datetime()).write_parquet( + data / "shard_0000.parquet" + ) + df = load_meds_events(tmp_path) + # Rows sorted by (subject_id, time). + assert df["subject_id"].to_list() == ["S1", "S1", "S2"] + + +def test_load_meds_events_filters_by_subject(tmp_path) -> None: + """subject_ids filter is honoured at scan time.""" + pl = pytest.importorskip("polars") + data = tmp_path / "data" + data.mkdir() + pl.DataFrame( + { + "subject_id": ["S1", "S2", "S3"], + "time": [ + "2020-01-01T00:00:00", + "2020-01-02T00:00:00", + "2020-01-03T00:00:00", + ], + "code": ["A", "B", "C"], + } + ).with_columns(pl.col("time").str.to_datetime()).write_parquet( + data / "shard_0000.parquet" + ) + df = load_meds_events(tmp_path, subject_ids=["S1", "S3"]) + assert sorted(df["subject_id"].to_list()) == ["S1", "S3"] + + +# --------------------------------------------------------------------------- +# Integration — only if meds_etl is importable +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _meds_etl_available, reason="meds_etl not installed") +def test_meds_etl_cli_registered() -> None: + """When meds_etl is installed, its console_scripts entry points register. + + We check via ``importlib.metadata`` rather than ``shutil.which`` so + the test passes under virtualenvs where the entry points exist but + the env's ``bin/`` isn't on the parent shell's ``$PATH`` (e.g. when + running pytest via an explicit interpreter path). + """ + import importlib.metadata as m + + ep_names: set[str] = set() + for dist in m.distributions(): + if (dist.metadata or {}).get("Name", "").lower() == "meds_etl": + ep_names.update(ep.name for ep in dist.entry_points) + assert "meds_etl_omop" in ep_names, ( + f"Expected meds_etl_omop console_script entry point; saw {ep_names}" + )