Skip to content

Commit

Permalink
feat: offer GZIP support for Metric IO
Browse files Browse the repository at this point in the history
  • Loading branch information
clintval committed Aug 22, 2023
1 parent f9c3774 commit 8d414b0
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
5 changes: 3 additions & 2 deletions fgpyo/util/metric.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@

import attr

from fgpyo import io
from fgpyo.util import inspect

MetricType = TypeVar("MetricType")
Expand Down Expand Up @@ -161,7 +162,7 @@ def read(cls, path: Path, ignore_extra_fields: bool = True) -> Iterator[Any]:
ignore_extra_fields: True to ignore any extra columns, False to raise an exception.
"""
parsers = cls._parsers()
with path.open("r") as reader:
with io.to_reader(path) as reader:
header: List[str] = reader.readline().rstrip("\r\n").split("\t")
# check the header
class_fields = set(cls.header())
Expand Down Expand Up @@ -234,7 +235,7 @@ def write(cls, path: Path, *values: MetricType) -> None:
path: path to the output file
values: zero or more metrics.
"""
with path.open("w") as writer:
with io.to_writer(path) as writer:
writer.write("\t".join(cls.header()))
writer.write("\n")
for value in values:
Expand Down
13 changes: 13 additions & 0 deletions fgpyo/util/tests/test_metric.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import enum
import gzip
from pathlib import Path
from typing import Any
from typing import Callable
Expand Down Expand Up @@ -157,6 +158,18 @@ def test_metrics_roundtrip(tmpdir: TmpDir) -> None:
assert len(metrics) == len(DUMMY_METRICS)
assert metrics == DUMMY_METRICS

def test_metric_roundtrip_gzip(tmp_path: Path) -> None:
path: Path = Path(tmp_path) / "metrics.txt.gz"

DummyMetric.write(path, *DUMMY_METRICS)

with gzip.open(path, "r") as handle:
handle.read(1) # Will raise an exception if not a GZIP file.

metrics: List[DummyMetric] = list(DummyMetric.read(path=path))

assert len(metrics) == len(DUMMY_METRICS)
assert metrics == DUMMY_METRICS

def test_metrics_read_extra_columns(tmpdir: TmpDir) -> None:
person = Person(name="Max", age=42)
Expand Down

0 comments on commit 8d414b0

Please sign in to comment.