testing code for fits strings - run test_roundtrips.py first, then test_read.py

test_read.py

import glob
import sys
import astropy
import astropy.table
import fitsio

IS_PY3 = tuple(sys.version_info)[0:3] >= (3, 0, 0)


def _read_fitsio(fname):
    return fitsio.read(fname)


def _read_astropy(fname):
    return astropy.table.Table.read(fname)


READ = {}
READ['fitsio'] = _read_fitsio
READ['astropy'] = _read_astropy

VERS = {}
VERS['fitsio'] = fitsio.__version__
VERS['astropy'] = astropy.__version__
codes = ['fitsio', 'astropy']

files = glob.glob('*.fits')
for fname in files:
    items = fname.split('_')

    if 'fitsio' in items[-2]:
        write_code = 'fitsio'
    else:
        write_code = 'astropy'
    write_vers = items[-2].replace(write_code, '')

    if write_code == 'astropy':
        nullterm = True
    elif write_code == 'fitsio' and '1.0.0' in write_vers:
        nullterm = True
    else:
        nullterm = False

    for read_code in codes:
        read_vers = VERS[read_code]

        d = READ[read_code](fname)
        if nullterm:
            val = 'a'
        else:
            val = 'a    '

        if IS_PY3 and read_code == 'fitsio' and '1.0.0' not in read_vers:
            val = val.encode('ascii')

        assert d['col'][0] == val

test_roundtrips.py

import sys
import numpy as np
import astropy
import astropy.table
import fitsio

IS_PY3 = tuple(sys.version_info)[0:3] >= (3, 0, 0)


def _read_fitsio(fname):
    return fitsio.read(fname)


def _write_fitsio(data, fname):
    fitsio.write(fname, data, clobber=True)


def _write_astropy(data, fname):
    astropy.table.Table(data).write(fname, overwrite=True)


def _read_astropy(fname):
    return astropy.table.Table.read(fname)


READ = {}
READ['fitsio'] = _read_fitsio
READ['astropy'] = _read_astropy

WRITE = {}
WRITE['fitsio'] = _write_fitsio
WRITE['astropy'] = _write_astropy

VERS = {}
VERS['fitsio'] = fitsio.__version__
VERS['astropy'] = astropy.__version__

strs = ['a', 'b', 'c']
dtypes = ['S5', 'U5']
codes = ['fitsio', 'astropy']
pys = [2, 3]

for read_code in codes:
    for write_code in codes:
        for dtype in dtypes:
            if not IS_PY3 and 'U' in dtype:
                continue

            data = np.zeros(3, dtype=[('col', dtype)])
            data['col'] = np.array(strs, dtype=dtype)

            if (write_code == 'fitsio' and
                    '1.0.0' not in fitsio.__version__ and
                    'U' in dtype):
                continue

            oname = 'test_data_dtype%s_%sv%s_py%d.fits' % (
                dtype, write_code, VERS[write_code],
                tuple(sys.version_info)[0])
            WRITE[write_code](data, oname)
            in_data = READ[read_code](oname)

            if read_code == 'fitsio':
                if IS_PY3 and '1.0.0' in fitsio.__version__:
                    assert in_data['col'].dtype.char == 'U'
                    assert in_data['col'][0] == 'a'
                else:
                    assert in_data['col'].dtype.char == 'S'

                    if IS_PY3:
                        if (write_code == 'astropy' or
                                '1.0.0' in fitsio.__version__):
                            pad = b''
                        else:
                            pad = b'    '
                        assert in_data['col'][0] == b'a' + pad
                    else:
                        if (write_code == 'astropy' or
                                '1.0.0' in fitsio.__version__):
                            pad = ''
                        else:
                            pad = '    '
                        assert in_data['col'][0] == 'a' + pad

            if read_code == 'astropy':
                assert in_data['col'].dtype.char == 'S'
                if write_code == 'astropy' or '1.0.0' in fitsio.__version__:
                    pad = ''
                else:
                    pad = '    '
                assert in_data['col'][0] == 'a' + pad