Support for hickle-ing astropy.Cosmology

[steven-murray]

It would be neat to be able to natively hickle astropy.cosmology.Cosmology subclass instances. I gave it a crack but my solution isn't working for me. I'm probably just misunderstanding something fundamental about how hickle works.

Here's my attempt:

def create_astropy_cosmology(py_obj, h_group, name, **kwargs):
    mapping = py_obj.to_format("mapping")
    c = mapping['cosmology']
    mapping['cosmology'] = (c.__module__, c.__name__)
    if 'astropy' not in c.__module__:
        warnings.warn("hickling a cosmology object that is not in astropy.cosmology. Written file will not be readable without loading the module defining the cosmology class.")
    return create_dictlike_dataset(mapping, h_group, name, **kwargs)

def load_astropy_cosmology(h_node,base_type,py_obj_type):
    asdict = hickle.load(h_node)
    cosmology_class = getattr(import_module(asdict['cosmology'][0]), asdict['cosmology'][1])
    return py_obj_type.from_format(asdict, cosmology=cosmology_class)

class_register += [
    [Cosmology, b'astropy_cosmology',  create_astropy_cosmology, load_astropy_cosmology],
    [FlatLambdaCDM, b'astropy_cosmology',  create_astropy_cosmology, load_astropy_cosmology],
]

@pytest.mark.parametrize('cosmo', [Planck15, WMAP1, Planck18])
def test_cosmology_roundtrip(cosmo):
    hdump(cosmo, f'cosmo-{cosmo.name}.hkl')
    new = hload(f'cosmo-{cosmo.name}.hkl')
    assert new == cosmo

However, when running the test, it always tells me DataRecoveredWarning: loader 'b'astropy_cosmology'' missing for 'FlatLambdaCDM' type object. Data recovered (data). Any idea why this is failing?

[telegraphic]

@steven-murray coming to this late, but I have no idea.

That error must be coming from astropy Cosmology (not hickle), so I guess the data isn't loading as expected. Did you find a solution -- is hickleable able to handle it?

[steven-murray]

Hey @telegraphic -- no I didn't find a solution. It's been a while and I can't remember the details, so I'm not sure there's more I can say about it right now. I'll be sure to keep it in mind though when it pops up again.