Skip to content

Commit 26f1d67

Browse files
authored
Merge pull request #23 from wkerzendorf/add_pdhdf_file_format
add initial pandas HDF fileformat
2 parents c1779e6 + 1e12df4 commit 26f1d67

File tree

6 files changed

+48
-1
lines changed

6 files changed

+48
-1
lines changed

CHANGES.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
0.6 (unreleased)
22
----------------
33

4+
- Add ability to compare to Pandas DataFrames and store them as HDF5 files [#23]
5+
46
0.5 (2022-01-12)
57
----------------
68

README.rst

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ in cases where the arrays are too large to conveniently hard-code them
1515
in the tests (e.g. ``np.testing.assert_allclose(x, [1, 2, 3])``).
1616

1717
The basic idea is that you can write a test that generates a Numpy array (or
18-
other related objects depending on the format). You can then either run the
18+
other related objects depending on the format, e.g. pandas DataFrame).
19+
You can then either run the
1920
tests in a mode to **generate** reference files from the arrays, or you can run
2021
the tests in **comparison** mode, which will compare the results of the tests to
2122
the reference ones within some tolerance.
@@ -25,6 +26,7 @@ At the moment, the supported file formats for the reference files are:
2526
- A plain text-based format (based on Numpy ``loadtxt`` output)
2627
- The FITS format (requires `astropy <http://www.astropy.org>`__). With this
2728
format, tests can return either a Numpy array for a FITS HDU object.
29+
- A pandas HDF5 format using the pandas HDFStore
2830

2931
For more information on how to write tests to do this, see the **Using**
3032
section below.

pytest_arraydiff/plugin.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,9 +137,43 @@ def write(filename, data, **kwargs):
137137
return np.savetxt(filename, data, **kwargs)
138138

139139

140+
class PDHDFDiff(BaseDiff):
141+
142+
extension = 'h5'
143+
144+
@staticmethod
145+
def read(filename):
146+
import pandas as pd
147+
return pd.read_hdf(filename)
148+
149+
@staticmethod
150+
def write(filename, data, **kwargs):
151+
import pandas as pd
152+
key = os.path.basename(filename).replace('.h5', '')
153+
return data.to_hdf(filename, key, **kwargs)
154+
155+
@classmethod
156+
def compare(cls, reference_file, test_file, atol=None, rtol=None):
157+
import pandas.testing as pdt
158+
import pandas as pd
159+
160+
ref_data = pd.read_hdf(reference_file)
161+
test_data = pd.read_hdf(test_file)
162+
try:
163+
pdt.assert_frame_equal(ref_data, test_data)
164+
except AssertionError as exc:
165+
message = "\n\na: {0}".format(test_file) + '\n'
166+
message += "b: {0}".format(reference_file) + '\n'
167+
message += exc.args[0]
168+
return False, message
169+
else:
170+
return True, ""
171+
172+
140173
FORMATS = {}
141174
FORMATS['fits'] = FITSDiff
142175
FORMATS['text'] = TextDiff
176+
FORMATS['pd_hdf'] = PDHDFDiff
143177

144178

145179
def _download_file(url):

setup.cfg

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ install_requires =
3636
[options.extras_require]
3737
test =
3838
astropy
39+
pandas
40+
tables
3941

4042
[options.entry_points]
4143
pytest11 =
6.87 KB
Binary file not shown.

tests/test_pytest_arraydiff.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,13 @@ def test_succeeds_func_text():
1818
return np.arange(3 * 5).reshape((3, 5))
1919

2020

21+
@pytest.mark.array_compare(file_format='pd_hdf', reference_dir=reference_dir)
22+
def test_succeeds_func_pdhdf():
23+
pd = pytest.importorskip('pandas')
24+
return pd.DataFrame(data=np.arange(20, dtype='int64'),
25+
columns=['test_data'])
26+
27+
2128
@pytest.mark.array_compare(file_format='fits', reference_dir=reference_dir)
2229
def test_succeeds_func_fits():
2330
return np.arange(3 * 5).reshape((3, 5)).astype(np.int64)

0 commit comments

Comments
 (0)