Source code for gwemlightcurves.KNModels.io.model

# -*- coding: utf-8 -*-
# Copyright (C) Duncan Macleod (2013)
#
# This file is part of GWpy.
#
# GWpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GWpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GWpy.  If not, see <http://www.gnu.org/licenses/>.

"""Fetch registration for database queries
"""

import re

from six import string_types

from astropy.io.registry import IORegistryError
from astropy.table import Table

_MODELS = {}

__author__ = 'Duncan Macleod <duncan.macleod@ligo.org>'


[docs]def register_model(data_format, data_class, function, force=False, usage=None): """Register a new method to EventTable.model() for a given format Parameters ---------- data_format : `str` name of the format to be registered data_class : `type` the class that the model returns function : `callable` the method to call from :meth:`EventTable.model` force : `bool`, optional overwrite existing registration for ``data_format`` if found, default: `False` """ key = (data_format, data_class) if key not in _MODELS or force: _MODELS[key] = (function, usage) else: raise IORegistryError("Fetcher for format '{0}' and class '{1}' " "has already been " "defined".format( data_format, data_class)) _update__doc__(data_class)
[docs]def get_model(data_format, data_class): """Return the :meth:`~KNTable.model` function for the given format Parameters ---------- data_format : `str` name of the format data_class : `type` the class that the model returns Raises ------ astropy.io.registry.IORegistryError if not registration is found matching ``data_format`` """ try: return _MODELS[(data_format, data_class)][0] except KeyError: formats = '\n'.join(_MODELS.keys()) raise IORegistryError("No model definer for format %r. " "The available formats are:\n%r" % (data_format, formats))
def _update__doc__(data_class): header = "The available named formats are:" model = data_class.model # if __doc__ isn't a string, bail-out now if not isinstance(model.__doc__, string_types): return # remove the old format list lines = model.__doc__.splitlines() try: pos = [i for i, line in enumerate(lines) if header in line][0] except IndexError: pass else: lines = lines[:pos] # work out the indentation matches = [re.search(r'(\S)', line) for line in lines[1:]] indent = min(match.start() for match in matches if match) # now re-write the format list formats = [] for fmt, cls in sorted(_MODELS, key=lambda x: x[0]): if cls is not data_class: continue usage = _MODELS[(fmt, cls)][1] formats.append(( fmt, '``model(%r, %s)``' % (fmt, usage))) format_str = Table(rows=formats, names=['Format', 'Basic usage']).pformat( max_lines=-1, max_width=80, align=('>', '<')) format_str[1] = format_str[1].replace('-', '=') format_str.insert(0, format_str[1]) format_str.append(format_str[0]) lines.extend([' ' * indent + line for line in [header, ''] + format_str]) # and overwrite the docstring try: model.__doc__ = '\n'.join(lines) except AttributeError: model.__func__.__doc__ = '\n'.join(lines)