<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst

from __future__ import print_function

import os

from ...utils import OrderedDict
from ...table import io_registry, Table
from . import HDUList, TableHDU, BinTableHDU
from . import open as fits_open


def is_fits(origin, args, kwargs):
    """
    Determine whether `origin` is a FITS file.

    Parameters
    ----------
    origin : str or readable file-like object
        Path or file object containing a potential FITS file.

    Returns
    -------
    is_fits : bool
        Returns `True` if the given file is a FITS file.
    """
    if isinstance(origin, basestring):
        if origin.lower().endswith(('.fits', '.fits.gz', '.fit', '.fit.gz')):
            return True
        else:
            return False
    else:
        return False

def read_table_fits(input, hdu_id=None):
    """
    Read a Table object from an FITS file

    Parameters
    ----------
    input : str or `~astropy.io.fits.hdu.table.TableHDU` or `~astropy.io.fits.hdu.table.BinTableHDU` or `~astropy.io.fits.hdu.hdulist.HDUList`
        If a string, the filename to read the table from. If a
        :class:`~astropy.io.fits.hdu.table.TableHDU` or
        :class:`~astropy.io.fits.hdu.table.BinTableHDU` or
        :class:`~astropy.io.fits.hdu.hdulist.HDUList`, the object to extract
        the table from.
    hdu_id : str, optional
        The HDU to read the table from
    """
    if isinstance(input, basestring):
        input = fits_open(input)

    # Parse all table objects
    tables = OrderedDict()
    if isinstance(input, HDUList):
        for ihdu, hdu in enumerate(input):
            if isinstance(hdu, (TableHDU, BinTableHDU)):
                tables[ihdu] = hdu

        if len(tables) > 1:
            if hdu_id is None:
                raise ValueError(
                    "Multiple tables found: HDU id should be set via "
                    "the hdu= argument. The available tables HDUs are " +
                    ', '.join([str(x) for x in tables.keys()]))
            else:
                if hdu_id in tables:
                    table = tables[hdu_id]
                else:
                    raise ValueError(
                        "No tables with hdu_id={0} found".format(hdu_id))
        elif len(tables) == 1:
            table = tables[tables.keys()[0]]
        else:
            raise ValueError("No table found")

    # Convert to an astropy.table.Table object
    return Table(table.data)


def write_table_fits(input, output, overwrite=False):
    """
    Write a Table object to a FITS file

    Parameters
    ----------
    input : Table
        The table to write out.
    output : str
        The filename to write the table to.
    overwrite : bool
        Whether to overwrite any existing file without warning.
    """

    # Check if output file already exists
    if isinstance(output, basestring) and os.path.exists(output):
        if overwrite:
            os.remove(output)
        else:
            raise IOError("File exists: {0}".format(output))

    # Create a new HDU object
    table_hdu = BinTableHDU(input._data)

    # Write out file
    table_hdu.writeto(output)


io_registry.register_reader('fits', read_table_fits)
io_registry.register_writer('fits', write_table_fits)
io_registry.register_identifier('fits', is_fits)
<commit_msg>Updated to latest registry API, and improved auto-detection of FITS format<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst

from __future__ import print_function

import os

from ...utils import OrderedDict
from .. import registry as io_registry
from ...table import Table

from . import HDUList, TableHDU, BinTableHDU
from . import open as fits_open


# FITS file signature as per RFC 4047
FITS_SIGNATURE = (b"\x53\x49\x4d\x50\x4c\x45\x20\x20\x3d\x20\x20\x20\x20\x20"
                  b"\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20"
                  b"\x20\x54")


def is_fits(origin, args, kwargs):
    """
    Determine whether `origin` is a FITS file.

    Parameters
    ----------
    origin : str or readable file-like object
        Path or file object containing a potential FITS file.

    Returns
    -------
    is_fits : bool
        Returns `True` if the given file is a FITS file.
    """
    if isinstance(args[0], basestring):
        if args[0].lower().endswith(('.fits', '.fits.gz', '.fit', '.fit.gz')):
            return True
        else:
            f = open(args[0], 'rb')
            sig = f.read(30)
            f.close()
            return sig == FITS_SIGNATURE
    elif hasattr(args[0], 'read'):
        pos = args[0].tell()
        sig = args[0].read(30)
        args[0].seek(pos)
        return sig == FITS_SIGNATURE
    elif isinstance(args[0], (HDUList, TableHDU, BinTableHDU)):
        return True
    else:
        return False


def read_table_fits(input, hdu_id=None):
    """
    Read a Table object from an FITS file

    Parameters
    ----------
    input : str or fileobj or `~astropy.io.fits.hdu.table.TableHDU` or `~astropy.io.fits.hdu.table.BinTableHDU` or `~astropy.io.fits.hdu.hdulist.HDUList`
        If a string, the filename to read the table from. If a file object, or
        a :class:`~astropy.io.fits.hdu.table.TableHDU` or
        :class:`~astropy.io.fits.hdu.table.BinTableHDU` or
        :class:`~astropy.io.fits.hdu.hdulist.HDUList` instance, the object to
        extract the table from.
    hdu_id : str, optional
        The HDU to read the table from
    """

    if isinstance(input, basestring):
        input = fits_open(input)

    # Parse all table objects
    tables = OrderedDict()
    if isinstance(input, HDUList):
        for ihdu, hdu in enumerate(input):
            if isinstance(hdu, (TableHDU, BinTableHDU)):
                tables[ihdu] = hdu

        if len(tables) > 1:
            if hdu_id is None:
                raise ValueError(
                    "Multiple tables found: HDU id should be set via "
                    "the hdu= argument. The available tables HDUs are " +
                    ', '.join([str(x) for x in tables.keys()]))
            else:
                if hdu_id in tables:
                    table = tables[hdu_id]
                else:
                    raise ValueError(
                        "No tables with hdu_id={0} found".format(hdu_id))
        elif len(tables) == 1:
            table = tables[tables.keys()[0]]
        else:
            raise ValueError("No table found")

    # Convert to an astropy.table.Table object
    t = Table(table.data)

    # TODO: deal properly with unsigned integers

    # TODO: read in metadata

    return t


def write_table_fits(input, output, overwrite=False):
    """
    Write a Table object to a FITS file

    Parameters
    ----------
    input : Table
        The table to write out.
    output : str
        The filename to write the table to.
    overwrite : bool
        Whether to overwrite any existing file without warning.
    """

    # Check if output file already exists
    if isinstance(output, basestring) and os.path.exists(output):
        if overwrite:
            os.remove(output)
        else:
            raise IOError("File exists: {0}".format(output))

    # Create a new HDU object
    table_hdu = BinTableHDU(input._data)

    # Write out file
    table_hdu.writeto(output)

    # TODO: write out metadata


io_registry.register_reader('fits', Table, read_table_fits)
io_registry.register_writer('fits', Table, write_table_fits)
io_registry.register_identifier('fits', Table, is_fits)

{'fraction_non_alphanumeric': 0.10606860158311346, 'fraction_numerical': 0.011213720316622692, 'mean_word_length': 2.082960553070354, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5907206', 'n_tokens_mistral': 2405, 'n_tokens_neox': 2266, 'n_words': 692}