/astropy/io/registry/core.py
Python | 404 lines | 373 code | 16 blank | 15 comment | 17 complexity | 5dc4d8a56721ede91c5b728ae74d3168 MD5 | raw file
- # Licensed under a 3-clause BSD style license - see LICENSE.rst
- import os
- import sys
- from collections import OrderedDict
- import numpy as np
- 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
- # -----------------------------------------------------------------------------
- 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 = OrderedDict()
- self._registries["read"] = dict(attr="_readers", column="Read")
- self._registries_order = ("read", "identify")
- # =========================================================================
- # Read methods
- 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 not (data_format, data_class) in self._readers or force:
- self._readers[(data_format, data_class)] = function, priority
- else:
- raise IORegistryError("Reader for format '{}' and class '{}' is "
- 'already defined'
- ''.format(data_format, data_class.__name__))
- if data_class not in self._delayed_docs_classes:
- self._update__doc__(data_class, 'read')
- 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("No reader defined for format '{}' and class '{}'"
- ''.format(data_format, data_class.__name__))
- if data_class not in self._delayed_docs_classes:
- self._update__doc__(data_class, 'read')
- 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]
- else:
- format_table_str = self._get_format_table_str(data_class, 'Read')
- raise IORegistryError(
- "No reader defined for format '{}' and class '{}'.\n\nThe "
- "available formats are:\n\n{}".format(
- data_format, data_class.__name__, format_table_str))
- 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)
- 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)
- except Exception:
- raise TypeError('could not convert reader output to {} '
- 'class.'.format(cls.__name__))
- finally:
- if ctx is not None:
- ctx.__exit__(*sys.exc_info())
- return data
- # -----------------------------------------------------------------------------
- class UnifiedOutputRegistry(_UnifiedIORegistryBase):
- """Write-only Registry.
- .. versionadded:: 5.0
- """
- def __init__(self):
- super().__init__()
- self._writers = OrderedDict()
- self._registries["write"] = dict(attr="_writers", column="Write")
- self._registries_order = ("write", "identify", )
- # =========================================================================
- # Write Methods
- 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:
- self._writers[(data_format, data_class)] = function, priority
- else:
- raise IORegistryError("Writer for format '{}' and class '{}' is "
- 'already defined'
- ''.format(data_format, data_class.__name__))
- if data_class not in self._delayed_docs_classes:
- self._update__doc__(data_class, 'write')
- 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("No writer defined for format '{}' and class '{}'"
- ''.format(data_format, data_class.__name__))
- if data_class not in self._delayed_docs_classes:
- self._update__doc__(data_class, 'write')
- 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]
- else:
- format_table_str = self._get_format_table_str(data_class, 'Write')
- raise IORegistryError(
- "No writer defined for format '{}' and class '{}'.\n\nThe "
- "available formats are:\n\n{}".format(
- data_format, data_class.__name__, format_table_str))
- 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)
- # -----------------------------------------------------------------------------
- class UnifiedIORegistry(UnifiedInputRegistry, UnifiedOutputRegistry):
- """Unified I/O Registry.
- .. versionadded:: 5.0
- """
- def __init__(self):
- super().__init__()
- self._registries_order = ("read", "write", "identify")
- 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)