# Path: sncosmo/registry.py
# def _get_registry(data_class):
# def register_loader(data_class, name, func, args=None,
#                     version=None, meta=None, force=False):
# def register(instance, name=None, data_class=None, force=False):
# def retrieve(data_class, name, version=None):


from io import StringIO
from numpy.testing import assert_allclose, assert_approx_equal
from sncosmo import registry
import numpy as np
import sncosmo
    """Effect with transmission 0 below cutoff wavelength, 1 above.
    Useful for testing behavior with redshift."""

    _param_names = ['stepwave']
    param_names_latex = [r'\lambda_{s}']

    def __init__(self, minwave, maxwave, stepwave=10000.):
        self._minwave = minwave
        self._maxwave = maxwave
        self._parameters = np.array([stepwave], dtype=np.float64)

    def propagate(self, wave, flux):
        return flux * (wave >= self._parameters[0])


class TestTimeSeriesSource:

    def setup_class(self):
        self.source = flatsource()

    def test_flux(self):
        for a in [1., 2.]:
            self.source.set(amplitude=a)
            assert_allclose(self.source.flux(1., 2000.), a)
            assert_allclose(self.source.flux(1., [2000.]), np.array([a]))
            assert_allclose(self.source.flux([1.], 2000.), np.array([[a]]))
            assert_allclose(self.source.flux([1.], [2000.]), np.array([[a]]))

    def test_getsource(self):
        # register the source & retrieve it from the registry

# Path: sncosmo/registry.py
# def _get_registry(data_class):
# def register_loader(data_class, name, func, args=None,
#                     version=None, meta=None, force=False):
# def register(instance, name=None, data_class=None, force=False):
# def retrieve(data_class, name, version=None):


# Path: sncosmo/tests/test_models.py
from io import StringIO
from numpy.testing import assert_allclose, assert_approx_equal
from sncosmo import registry
import numpy as np
import sncosmo

    """Effect with transmission 0 below cutoff wavelength, 1 above.
    Useful for testing behavior with redshift."""

    _param_names = ['stepwave']
    param_names_latex = [r'\lambda_{s}']

    def __init__(self, minwave, maxwave, stepwave=10000.):
        self._minwave = minwave
        self._maxwave = maxwave
        self._parameters = np.array([stepwave], dtype=np.float64)

    def propagate(self, wave, flux):
        return flux * (wave >= self._parameters[0])


class TestTimeSeriesSource:

    def setup_class(self):
        self.source = flatsource()

    def test_flux(self):
        for a in [1., 2.]:
            self.source.set(amplitude=a)
            assert_allclose(self.source.flux(1., 2000.), a)
            assert_allclose(self.source.flux(1., [2000.]), np.array([a]))
            assert_allclose(self.source.flux([1.], 2000.), np.array([[a]]))
            assert_allclose(self.source.flux([1.], [2000.]), np.array([[a]]))

    def test_getsource(self):
        # register the source & retrieve it from the registry

        registry.register(self.source, name="testsource",