<commit_before>import gc
import sys
import pathlib
import warnings

import pytest
import numpy as np
from numpy.testing import assert_allclose, assert_array_equal

from astropy.io.fits.column import (_parse_tdisp_format, _fortran_to_python_format,
                                    python_to_tdisp)

from astropy.io.fits import HDUList, PrimaryHDU, BinTableHDU, ImageHDU, table_to_hdu

from astropy.io import fits

from astropy import units as u
from astropy.table import Table, QTable, NdarrayMixin, Column
from astropy.table.table_helpers import simple_table
from astropy.units.format.fits import UnitScaleError
from astropy.utils.data import get_pkg_data_filename
from astropy.utils.exceptions import (AstropyUserWarning,
                                      AstropyDeprecationWarning)

from astropy.coordinates import (SkyCoord, Latitude, Longitude, Angle, EarthLocation,
                                 SphericalRepresentation, CartesianRepresentation,
                                 SphericalCosLatDifferential)
from astropy.time import Time, TimeDelta
from astropy.units.quantity import QuantityInfo
from astropy.utils.compat.optional_deps import HAS_YAML  # noqa


def equal_data(a, b):
    for name in a.dtype.names:
        if not np.all(a[name] == b[name]):
            return False
    return True


class TestSingleTable:

    def setup_class(self):
        self.data = np.array(list(zip([1, 2, 3, 4],
                                      ['a', 'b', 'c', 'd'],
                                      [2.3, 4.5, 6.7, 8.9])),
                             dtype=[('a', int), ('b', 'U1'), ('c', float)])

    def test_simple(self, tmpdir):
        filename = str(tmpdir.join('test_simple.fts'))
        t1 = Table(self.data)
        t1.write(filename, overwrite=True)
        t2 = Table.read(filename)
        assert equal_data(t1, t2)

    def test_simple_pathlib(self, tmpdir):
        filename = pathlib.Path(str(tmpdir.join('test_simple.fit')))
        t1 = Table(self.data)
        t1.write(filename, overwrite=True)
        t2 = Table.read(filename)
        assert equal_data(t1, t2)

    def test_simple_meta(self, tmpdir):
        filename = str(tmpdir.join('test_simple.fits'))
        t1 = Table(self.data)
        t1.meta['A'] = 1
        t1.meta['B'] = 2.3
        t1.meta['C'] = 'spam'
        t1.meta['comments'] = ['this', 'is', 'a', 'long', 'comment']
        t1.meta['HISTORY'] = ['first', 'second', 'third']
        t1.write(filename, overwrite=True)
        t2 = Table.read(filename)
        assert equal_data(t1, t2)
        for key in t1.meta:
            if isinstance(t1.meta, list):
                for i in range(len(t1.meta[key])):
                    assert t1.meta[key][i] == t2.meta[key][i]
            else:
                assert t1.meta[key] == t2.meta[key]

    def test_simple_meta_conflicting(self, tmpdir):
        filename = str(tmpdir.join('test_simple.fits'))
        t1 = Table(self.data)
        t1.meta['ttype1'] = 'spam'
        with pytest.warns(AstropyUserWarning, match='Meta-data keyword ttype1 '
                          'will be ignored since it conflicts with a FITS '
                          'reserved keyword') as w:
            t1.write(filename, overwrite=True)
        assert len(w) == 1

    def test_simple_noextension(self, tmpdir):
        """
        Test that file type is recognized without extension
        """
        filename = str(tmpdir.join('test_simple'))
        t1 = Table(self.data)
        t1.write(filename, overwrite=True, format='fits')
        t2 = Table.read(filename)
        assert equal_data(t1, t2)

    @pytest.mark.parametrize('table_type', (Table, QTable))
    def test_with_units(self, table_type, tmpdir):
        filename = str(tmpdir.join('test_with_units.fits'))
        t1 = table_type(self.data)
        t1['a'].unit = u.m
        t1['c'].unit = u.km / u.s
        t1.write(filename, overwrite=True)
        t2 = table_type.read(filename)
        assert equal_data(t1, t2)
        assert t2['a'].unit == u.m
        assert t2['c'].unit == u.km / u.s

    @pytest.mark.skipif('not HAS_YAML')
    def test_with_custom_units_qtable(self, tmpdir):
        # Test only for QTable - for Table's Column, new units are dropped
        # (as is checked in test_write_drop_nonstandard_units).
        filename = str(tmpdir.join('test_with_units.fits'))
        unit = u.def_unit('bandpass_sol_lum')
        t = QTable()
        t['l'] = np.ones(5) * unit
        with pytest.warns(AstropyUserWarning) as w:
            t.write(filename, overwrite=True)
        assert len(w) == 1
        assert 'bandpass_sol_lum' in str(w[0].message)
        # Just reading back, the data is fine but the unit is not recognized.
        with pytest.warns(u.UnitsWarning, match="'bandpass_sol_lum' did not parse") as w:
            t2 = QTable.read(filename)
        assert len(w) == 1
        assert isinstance(t2['l'].unit, u.UnrecognizedUnit)
        assert str(t2['l'].unit) == 'bandpass_sol_lum'
        assert np.all(t2['l'].value == t['l'].value)

        # But if we enable the unit, it should be recognized.
        with u.add_enabled_units(unit):
            t3 = QTable.read(filename)
            assert t3['l'].unit is unit
            assert equal_data(t3, t)

            # Regression check for #8897; write used to fail when a custom
            # unit was enabled.
            with pytest.warns(AstropyUserWarning):
                t3.write(filename, overwrite=True)

        # It should also be possible to read the file in using a unit alias,
        # even to a unit that may not be the same.
        with u.set_enabled_aliases({'bandpass_sol_lum': u.Lsun}):
            t3 = QTable.read(filename)
            assert t3['l'].unit is u.Lsun

    @pytest.mark.parametrize('table_type', (Table, QTable))
    def test_with_format(self, table_type, tmpdir):
        filename = str(tmpdir.join('test_with_format.fits'))
        t1 = table_type(self.data)
        t1['a'].format = '{:5d}'
        t1['b'].format = '{:>20}'
        t1['c'].format = '{:6.2f}'
        t1.write(filename, overwrite=True)
        t2 = table_type.read(filename)
        assert equal_data(t1, t2)
        assert t2['a'].format == '{:5d}'
        assert t2['b'].format == '{:>20}'
        assert t2['c'].format == '{:6.2f}'

    def test_masked(self, tmpdir):
        filename = str(tmpdir.join('test_masked.fits'))
        t1 = Table(self.data, masked=True)
        t1.mask['a'] = [1, 0, 1, 0]
        t1.mask['b'] = [1, 0, 0, 1]
        t1.mask['c'] = [0, 1, 1, 0]
        t1.write(filename, overwrite=True)<commit_msg>Add test that unit aliases work in plain FITS files<commit_after>import gc
import sys
import pathlib
import warnings

import pytest
import numpy as np
from numpy.testing import assert_allclose, assert_array_equal

from astropy.io.fits.column import (_parse_tdisp_format, _fortran_to_python_format,
                                    python_to_tdisp)

from astropy.io.fits import HDUList, PrimaryHDU, BinTableHDU, ImageHDU, table_to_hdu

from astropy.io import fits

from astropy import units as u
from astropy.table import Table, QTable, NdarrayMixin, Column
from astropy.table.table_helpers import simple_table
from astropy.units.format.fits import UnitScaleError
from astropy.utils.data import get_pkg_data_filename
from astropy.utils.exceptions import (AstropyUserWarning,
                                      AstropyDeprecationWarning)

from astropy.coordinates import (SkyCoord, Latitude, Longitude, Angle, EarthLocation,
                                 SphericalRepresentation, CartesianRepresentation,
                                 SphericalCosLatDifferential)
from astropy.time import Time, TimeDelta
from astropy.units.quantity import QuantityInfo
from astropy.utils.compat.optional_deps import HAS_YAML  # noqa


def equal_data(a, b):
    for name in a.dtype.names:
        if not np.all(a[name] == b[name]):
            return False
    return True


class TestSingleTable:

    def setup_class(self):
        self.data = np.array(list(zip([1, 2, 3, 4],
                                      ['a', 'b', 'c', 'd'],
                                      [2.3, 4.5, 6.7, 8.9])),
                             dtype=[('a', int), ('b', 'U1'), ('c', float)])

    def test_simple(self, tmpdir):
        filename = str(tmpdir.join('test_simple.fts'))
        t1 = Table(self.data)
        t1.write(filename, overwrite=True)
        t2 = Table.read(filename)
        assert equal_data(t1, t2)

    def test_simple_pathlib(self, tmpdir):
        filename = pathlib.Path(str(tmpdir.join('test_simple.fit')))
        t1 = Table(self.data)
        t1.write(filename, overwrite=True)
        t2 = Table.read(filename)
        assert equal_data(t1, t2)

    def test_simple_meta(self, tmpdir):
        filename = str(tmpdir.join('test_simple.fits'))
        t1 = Table(self.data)
        t1.meta['A'] = 1
        t1.meta['B'] = 2.3
        t1.meta['C'] = 'spam'
        t1.meta['comments'] = ['this', 'is', 'a', 'long', 'comment']
        t1.meta['HISTORY'] = ['first', 'second', 'third']
        t1.write(filename, overwrite=True)
        t2 = Table.read(filename)
        assert equal_data(t1, t2)
        for key in t1.meta:
            if isinstance(t1.meta, list):
                for i in range(len(t1.meta[key])):
                    assert t1.meta[key][i] == t2.meta[key][i]
            else:
                assert t1.meta[key] == t2.meta[key]

    def test_simple_meta_conflicting(self, tmpdir):
        filename = str(tmpdir.join('test_simple.fits'))
        t1 = Table(self.data)
        t1.meta['ttype1'] = 'spam'
        with pytest.warns(AstropyUserWarning, match='Meta-data keyword ttype1 '
                          'will be ignored since it conflicts with a FITS '
                          'reserved keyword') as w:
            t1.write(filename, overwrite=True)
        assert len(w) == 1

    def test_simple_noextension(self, tmpdir):
        """
        Test that file type is recognized without extension
        """
        filename = str(tmpdir.join('test_simple'))
        t1 = Table(self.data)
        t1.write(filename, overwrite=True, format='fits')
        t2 = Table.read(filename)
        assert equal_data(t1, t2)

    @pytest.mark.parametrize('table_type', (Table, QTable))
    def test_with_units(self, table_type, tmpdir):
        filename = str(tmpdir.join('test_with_units.fits'))
        t1 = table_type(self.data)
        t1['a'].unit = u.m
        t1['c'].unit = u.km / u.s
        t1.write(filename, overwrite=True)
        t2 = table_type.read(filename)
        assert equal_data(t1, t2)
        assert t2['a'].unit == u.m
        assert t2['c'].unit == u.km / u.s

    @pytest.mark.skipif('not HAS_YAML')
    def test_with_custom_units_qtable(self, tmpdir):
        # Test only for QTable - for Table's Column, new units are dropped
        # (as is checked in test_write_drop_nonstandard_units).
        filename = str(tmpdir.join('test_with_units.fits'))
        unit = u.def_unit('bandpass_sol_lum')
        t = QTable()
        t['l'] = np.ones(5) * unit
        with pytest.warns(AstropyUserWarning) as w:
            t.write(filename, overwrite=True)
        assert len(w) == 1
        assert 'bandpass_sol_lum' in str(w[0].message)
        # Just reading back, the data is fine but the unit is not recognized.
        with pytest.warns(u.UnitsWarning, match="'bandpass_sol_lum' did not parse") as w:
            t2 = QTable.read(filename)
        assert len(w) == 1
        assert isinstance(t2['l'].unit, u.UnrecognizedUnit)
        assert str(t2['l'].unit) == 'bandpass_sol_lum'
        assert np.all(t2['l'].value == t['l'].value)

        # But if we enable the unit, it should be recognized.
        with u.add_enabled_units(unit):
            t3 = QTable.read(filename)
            assert t3['l'].unit is unit
            assert equal_data(t3, t)

            # Regression check for #8897; write used to fail when a custom
            # unit was enabled.
            with pytest.warns(AstropyUserWarning):
                t3.write(filename, overwrite=True)

        # It should also be possible to read the file in using a unit alias,
        # even to a unit that may not be the same.
        with u.set_enabled_aliases({'bandpass_sol_lum': u.Lsun}):
            t3 = QTable.read(filename)
            assert t3['l'].unit is u.Lsun

    @pytest.mark.parametrize('table_type', (Table, QTable))
    def test_read_with_unit_aliases(self, table_type):
        hdu = BinTableHDU(self.data)
        hdu.columns[0].unit = 'Angstroms'
        hdu.columns[2].unit = 'ergs/(cm.s.Angstroms)'
        with u.set_enabled_aliases(dict(Angstroms=u.AA, ergs=u.erg)):
            t = table_type.read(hdu)
        assert t['a'].unit == u.AA
        assert t['c'].unit == u.erg/(u.cm*u.s*u.AA)

    @pytest.mark.parametrize('table_type', (Table, QTable))
    def test_with_format(self, table_type, tmpdir):
        filename = str(tmpdir.join('test_with_format.fits'))
        t1 = table_type(self.data)
        t1['a'].format = '{:5d}'
        t1['b'].format = '{:>20}'
        t1['c'].format = '{:6.2f}'
        t1.write(filename, overwrite=True)
        t2 = table_type.read(filename)
        assert equal_data(t1, t2)
        assert t2['a'].format == '{:5d}'
        assert t2['b'].format == '{:>20}'
        assert t2['c'].format == '{:6.2f}'

    def test_masked(self, tmpdir):
        filename = str(tmpdir.join('test_masked.fits'))
        t1 = Table(self.data, masked=True)
        t1.mask['a'] = [1, 0, 1, 0]
        t1.mask['b'] = [1, 0, 0, 1]
        t1.mask['c'] = [0, 1, 1, 0]
        t1.write(filename, overwrite=True)
{'fraction_non_alphanumeric': 0.12724079900490232, 'fraction_numerical': 0.0175605473037243, 'mean_word_length': 2.3206997084548107, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4136040', 'n_tokens_mistral': 4595, 'n_tokens_neox': 4440, 'n_words': 1029}