Source code for astropy.io.registry.core

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

import os
import sys

from .base import IORegistryError, _UnifiedIORegistryBase

__all__ = ["UnifiedIORegistry", "UnifiedInputRegistry", "UnifiedOutputRegistry"]


PATH_TYPES = (str, os.PathLike)  # TODO! include bytes


def _expand_user_in_args(args):
    # Conservatively attempt to apply `os.path.expanduser` to the first
    # argument, which can be either a path or the contents of a table.
    if len(args) and isinstance(args[0], PATH_TYPES):
        ex_user = os.path.expanduser(args[0])
        if ex_user != args[0] and os.path.exists(os.path.dirname(ex_user)):
            args = (ex_user,) + args[1:]
    return args


# -----------------------------------------------------------------------------


[docs] class UnifiedInputRegistry(_UnifiedIORegistryBase): """Read-only Unified Registry. .. versionadded:: 5.0 Examples -------- First let's start by creating a read-only registry. .. code-block:: python >>> from astropy.io.registry import UnifiedInputRegistry >>> read_reg = UnifiedInputRegistry() There is nothing in this registry. Let's make a reader for the :class:`~astropy.table.Table` class:: from astropy.table import Table def my_table_reader(filename, some_option=1): # Read in the table by any means necessary return table # should be an instance of Table Such a function can then be registered with the I/O registry:: read_reg.register_reader('my-table-format', Table, my_table_reader) Note that we CANNOT then read in a table with:: d = Table.read('my_table_file.mtf', format='my-table-format') Why? because ``Table.read`` uses Astropy's default global registry and this is a separate registry. Instead we can read by the read method on the registry:: d = read_reg.read(Table, 'my_table_file.mtf', format='my-table-format') """ def __init__(self): super().__init__() # set _identifiers self._readers = {} self._registries["read"] = {"attr": "_readers", "column": "Read"} self._registries_order = ("read", "identify") # ========================================================================= # Read methods
[docs] def register_reader( self, data_format, data_class, function, force=False, priority=0 ): """ Register a reader function. Parameters ---------- data_format : str The data format identifier. This is the string that will be used to specify the data type when reading. data_class : class The class of the object that the reader produces. function : function The function to read in a data object. force : bool, optional Whether to override any existing function if already present. Default is ``False``. priority : int, optional The priority of the reader, used to compare possible formats when trying to determine the best reader to use. Higher priorities are preferred over lower priorities, with the default priority being 0 (negative numbers are allowed though). """ if (data_format, data_class) not in self._readers or force: self._readers[(data_format, data_class)] = function, priority else: raise IORegistryError( f"Reader for format '{data_format}' and class '{data_class.__name__}'" " is already defined" ) if data_class not in self._delayed_docs_classes: self._update__doc__(data_class, "read")
[docs] def unregister_reader(self, data_format, data_class): """ Unregister a reader function. Parameters ---------- data_format : str The data format identifier. data_class : class The class of the object that the reader produces. """ if (data_format, data_class) in self._readers: self._readers.pop((data_format, data_class)) else: raise IORegistryError( f"No reader defined for format '{data_format}' and class" f" '{data_class.__name__}'" ) if data_class not in self._delayed_docs_classes: self._update__doc__(data_class, "read")
[docs] def get_reader(self, data_format, data_class): """Get reader for ``data_format``. Parameters ---------- data_format : str The data format identifier. This is the string that is used to specify the data type when reading/writing. data_class : class The class of the object that can be written. Returns ------- reader : callable The registered reader function for this format and class. """ readers = [(fmt, cls) for fmt, cls in self._readers if fmt == data_format] for reader_format, reader_class in readers: if self._is_best_match(data_class, reader_class, readers): return self._readers[(reader_format, reader_class)][0] format_table_str = self._get_format_table_str(data_class, "Read") raise IORegistryError( f"No reader defined for format '{data_format}' and class" f" '{data_class.__name__}'.\n\nThe available formats" f" are:\n\n{format_table_str}" )
[docs] def read(self, cls, *args, format=None, cache=False, **kwargs): """ Read in data. Parameters ---------- cls : class *args The arguments passed to this method depend on the format. format : str or None cache : bool Whether to cache the results of reading in the data. **kwargs The arguments passed to this method depend on the format. Returns ------- object or None The output of the registered reader. """ ctx = None try: # Expand a tilde-prefixed path if present in args[0] args = _expand_user_in_args(args) if format is None: path = None fileobj = None if len(args): if isinstance(args[0], PATH_TYPES) and not os.path.isdir(args[0]): from astropy.utils.data import get_readable_fileobj # path might be a os.PathLike object if isinstance(args[0], os.PathLike): args = (os.fspath(args[0]),) + args[1:] path = args[0] try: ctx = get_readable_fileobj( args[0], encoding="binary", cache=cache ) fileobj = ctx.__enter__() except OSError: raise except Exception: fileobj = None else: args = [fileobj] + list(args[1:]) elif hasattr(args[0], "read"): path = None fileobj = args[0] format = self._get_valid_format( "read", cls, path, fileobj, args, kwargs ) # We need to keep track the original path in case it uses a # relative path to the parquet binary if format == "votable": kwargs.update({"filename": path}) reader = self.get_reader(format, cls) data = reader(*args, **kwargs) if not isinstance(data, cls): # User has read with a subclass where only the parent class is # registered. This returns the parent class, so try coercing # to desired subclass. try: data = cls(data, copy=False) except Exception: raise TypeError( f"could not convert reader output to {cls.__name__} class." ) finally: if ctx is not None: ctx.__exit__(*sys.exc_info()) return data
# -----------------------------------------------------------------------------
[docs] class UnifiedOutputRegistry(_UnifiedIORegistryBase): """Write-only Registry. .. versionadded:: 5.0 """ def __init__(self): super().__init__() self._writers = {} self._registries["write"] = {"attr": "_writers", "column": "Write"} self._registries_order = ("write", "identify") # ========================================================================= # Write Methods
[docs] def register_writer( self, data_format, data_class, function, force=False, priority=0 ): """ Register a table writer function. Parameters ---------- data_format : str The data format identifier. This is the string that will be used to specify the data type when writing. data_class : class The class of the object that can be written. function : function The function to write out a data object. force : bool, optional Whether to override any existing function if already present. Default is ``False``. priority : int, optional The priority of the writer, used to compare possible formats when trying to determine the best writer to use. Higher priorities are preferred over lower priorities, with the default priority being 0 (negative numbers are allowed though). """ if not (data_format, data_class) in self._writers or force: # noqa: E713 self._writers[(data_format, data_class)] = function, priority else: raise IORegistryError( f"Writer for format '{data_format}' and class '{data_class.__name__}'" " is already defined" ) if data_class not in self._delayed_docs_classes: self._update__doc__(data_class, "write")
[docs] def unregister_writer(self, data_format, data_class): """ Unregister a writer function. Parameters ---------- data_format : str The data format identifier. data_class : class The class of the object that can be written. """ if (data_format, data_class) in self._writers: self._writers.pop((data_format, data_class)) else: raise IORegistryError( f"No writer defined for format '{data_format}' and class" f" '{data_class.__name__}'" ) if data_class not in self._delayed_docs_classes: self._update__doc__(data_class, "write")
[docs] def get_writer(self, data_format, data_class): """Get writer for ``data_format``. Parameters ---------- data_format : str The data format identifier. This is the string that is used to specify the data type when reading/writing. data_class : class The class of the object that can be written. Returns ------- writer : callable The registered writer function for this format and class. """ writers = [(fmt, cls) for fmt, cls in self._writers if fmt == data_format] for writer_format, writer_class in writers: if self._is_best_match(data_class, writer_class, writers): return self._writers[(writer_format, writer_class)][0] format_table_str = self._get_format_table_str(data_class, "Write") raise IORegistryError( f"No writer defined for format '{data_format}' and class" f" '{data_class.__name__}'.\n\nThe available formats" f" are:\n\n{format_table_str}" )
[docs] def write(self, data, *args, format=None, **kwargs): """ Write out data. Parameters ---------- data : object The data to write. *args The arguments passed to this method depend on the format. format : str or None **kwargs The arguments passed to this method depend on the format. Returns ------- object or None The output of the registered writer. Most often `None`. .. versionadded:: 4.3 """ # Expand a tilde-prefixed path if present in args[0] args = _expand_user_in_args(args) if format is None: path = None fileobj = None if len(args): if isinstance(args[0], PATH_TYPES): # path might be a os.PathLike object if isinstance(args[0], os.PathLike): args = (os.fspath(args[0]),) + args[1:] path = args[0] fileobj = None elif hasattr(args[0], "read"): path = None fileobj = args[0] format = self._get_valid_format( "write", data.__class__, path, fileobj, args, kwargs ) writer = self.get_writer(format, data.__class__) return writer(data, *args, **kwargs)
# -----------------------------------------------------------------------------
[docs] class UnifiedIORegistry(UnifiedInputRegistry, UnifiedOutputRegistry): """Unified I/O Registry. .. versionadded:: 5.0 """ def __init__(self): super().__init__() self._registries_order = ("read", "write", "identify")
[docs] def get_formats(self, data_class=None, readwrite=None): """ Get the list of registered I/O formats as a `~astropy.table.Table`. Parameters ---------- data_class : class, optional Filter readers/writer to match data class (default = all classes). readwrite : str or None, optional Search only for readers (``"Read"``) or writers (``"Write"``). If None search for both. Default is None. .. versionadded:: 1.3 Returns ------- format_table : :class:`~astropy.table.Table` Table of available I/O formats. """ return super().get_formats(data_class, readwrite)