1# Licensed under a 3-clause BSD style license - see LICENSE.rst
2
3
4import os
5
6
7from . import parse, from_table
8from .tree import VOTableFile, Table as VOTable
9from astropy.io import registry as io_registry
10from astropy.table import Table
11from astropy.table.column import BaseColumn
12from astropy.units import Quantity
13from astropy.utils.misc import NOT_OVERWRITING_MSG
14
15
16def is_votable(origin, filepath, fileobj, *args, **kwargs):
17    """
18    Reads the header of a file to determine if it is a VOTable file.
19
20    Parameters
21    ----------
22    origin : str or readable file-like
23        Path or file object containing a VOTABLE_ xml file.
24
25    Returns
26    -------
27    is_votable : bool
28        Returns `True` if the given file is a VOTable file.
29    """
30    from . import is_votable
31    if origin == 'read':
32        if fileobj is not None:
33            try:
34                result = is_votable(fileobj)
35            finally:
36                fileobj.seek(0)
37            return result
38        elif filepath is not None:
39            return is_votable(filepath)
40        elif isinstance(args[0], (VOTableFile, VOTable)):
41            return True
42        else:
43            return False
44    else:
45        return False
46
47
48def read_table_votable(input, table_id=None, use_names_over_ids=False,
49                       verify=None, **kwargs):
50    """
51    Read a Table object from an VO table file
52
53    Parameters
54    ----------
55    input : str or `~astropy.io.votable.tree.VOTableFile` or `~astropy.io.votable.tree.Table`
56        If a string, the filename to read the table from. If a
57        :class:`~astropy.io.votable.tree.VOTableFile` or
58        :class:`~astropy.io.votable.tree.Table` object, the object to extract
59        the table from.
60
61    table_id : str or int, optional
62        The table to read in.  If a `str`, it is an ID corresponding
63        to the ID of the table in the file (not all VOTable files
64        assign IDs to their tables).  If an `int`, it is the index of
65        the table in the file, starting at 0.
66
67    use_names_over_ids : bool, optional
68        When `True` use the ``name`` attributes of columns as the names
69        of columns in the `~astropy.table.Table` instance.  Since names
70        are not guaranteed to be unique, this may cause some columns
71        to be renamed by appending numbers to the end.  Otherwise
72        (default), use the ID attributes as the column names.
73
74    verify : {'ignore', 'warn', 'exception'}, optional
75        When ``'exception'``, raise an error when the file violates the spec,
76        otherwise either issue a warning (``'warn'``) or silently continue
77        (``'ignore'``). Warnings may be controlled using the standard Python
78        mechanisms.  See the `warnings` module in the Python standard library
79        for more information. When not provided, uses the configuration setting
80        ``astropy.io.votable.verify``, which defaults to ``'ignore'``.
81
82    **kwargs
83        Additional keyword arguments are passed on to
84        :func:`astropy.io.votable.table.parse`.
85    """
86    if not isinstance(input, (VOTableFile, VOTable)):
87        input = parse(input, table_id=table_id, verify=verify, **kwargs)
88
89    # Parse all table objects
90    table_id_mapping = dict()
91    tables = []
92    if isinstance(input, VOTableFile):
93        for table in input.iter_tables():
94            if table.ID is not None:
95                table_id_mapping[table.ID] = table
96            tables.append(table)
97
98        if len(tables) > 1:
99            if table_id is None:
100                raise ValueError(
101                    "Multiple tables found: table id should be set via "
102                    "the table_id= argument. The available tables are {}, "
103                    'or integers less than {}.'.format(
104                        ', '.join(table_id_mapping.keys()), len(tables)))
105            elif isinstance(table_id, str):
106                if table_id in table_id_mapping:
107                    table = table_id_mapping[table_id]
108                else:
109                    raise ValueError(
110                        f"No tables with id={table_id} found")
111            elif isinstance(table_id, int):
112                if table_id < len(tables):
113                    table = tables[table_id]
114                else:
115                    raise IndexError(
116                        "Table index {} is out of range. "
117                        "{} tables found".format(
118                            table_id, len(tables)))
119        elif len(tables) == 1:
120            table = tables[0]
121        else:
122            raise ValueError("No table found")
123    elif isinstance(input, VOTable):
124        table = input
125
126    # Convert to an astropy.table.Table object
127    return table.to_table(use_names_over_ids=use_names_over_ids)
128
129
130def write_table_votable(input, output, table_id=None, overwrite=False,
131                        tabledata_format=None):
132    """
133    Write a Table object to an VO table file
134
135    Parameters
136    ----------
137    input : Table
138        The table to write out.
139
140    output : str
141        The filename to write the table to.
142
143    table_id : str, optional
144        The table ID to use. If this is not specified, the 'ID' keyword in the
145        ``meta`` object of the table will be used.
146
147    overwrite : bool, optional
148        Whether to overwrite any existing file without warning.
149
150    tabledata_format : str, optional
151        The format of table data to write.  Must be one of ``tabledata``
152        (text representation), ``binary`` or ``binary2``.  Default is
153        ``tabledata``.  See :ref:`astropy:votable-serialization`.
154    """
155
156    # Only those columns which are instances of BaseColumn or Quantity can be written
157    unsupported_cols = input.columns.not_isinstance((BaseColumn, Quantity))
158    if unsupported_cols:
159        unsupported_names = [col.info.name for col in unsupported_cols]
160        raise ValueError('cannot write table with mixin column(s) {} to VOTable'
161                         .format(unsupported_names))
162
163    # Check if output file already exists
164    if isinstance(output, str) and os.path.exists(output):
165        if overwrite:
166            os.remove(output)
167        else:
168            raise OSError(NOT_OVERWRITING_MSG.format(output))
169
170    # Create a new VOTable file
171    table_file = from_table(input, table_id=table_id)
172
173    # Write out file
174    table_file.to_xml(output, tabledata_format=tabledata_format)
175
176
177io_registry.register_reader('votable', Table, read_table_votable)
178io_registry.register_writer('votable', Table, write_table_votable)
179io_registry.register_identifier('votable', Table, is_votable)
180