File size: 1,598 Bytes
35cdf53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
"""Module that provides an abstraction for accessing test data."""
import os
import pathlib
from typing import Literal, overload
from absl.testing import absltest
class Data:
"""Provides an abstraction for accessing test data."""
def __init__(self, data_dir: os.PathLike[str] | str):
"""Initiailizes data wrapper, providing users with high level data access.
Args:
data_dir: Directory containing test data.
"""
self._data_dir = pathlib.Path(data_dir)
def path(self, data_name: str | os.PathLike[str] | None = None) -> str:
"""Returns the path to a given test data.
Args:
data_name: the name of the test data file relative to data_dir. If not
set, this will return the absolute path to the data directory.
"""
data_dir_path = (
pathlib.Path(absltest.get_default_test_srcdir()) / self._data_dir
)
if data_name:
return str(data_dir_path / data_name)
return str(data_dir_path)
@overload
def load(
self, data_name: str | os.PathLike[str], mode: Literal['rt'] = 'rt'
) -> str:
...
@overload
def load(
self, data_name: str | os.PathLike[str], mode: Literal['rb'] = 'rb'
) -> bytes:
...
def load(
self, data_name: str | os.PathLike[str], mode: str = 'rt'
) -> str | bytes:
"""Returns the contents of a given test data.
Args:
data_name: the name of the test data file relative to data_dir.
mode: the mode in which to read the data file. Defaults to text ('rt').
"""
with open(self.path(data_name), mode=mode) as f:
return f.read()
|