Make astropy.utils.masked.Masked work with dill

[Jerry-Ma]

Description

It appears that the following code to pickle a MaskedQuantity object does not work because the __setstate__ does not handle non-bound case:

import astropy.units as u
import dill as pickle
import numpy as np
from astropy.utils.masked import Masked

a = np.arange(3) << u.m
m = np.array([True, False, False])

a_masked = Masked(a, mask=m)

print(a_masked)


print(pickle.loads(pickle.dumps(a_masked)))

The fix seems to be simple: add logic to handle the uninitialized _attrs attribute on the astropy.utils.data_info.DataInfo class:

    def __getstate__(self):
        return getattr(self, "_attrs", None)

    def __setstate__(self, state):
        if state is not None:
            self._attrs = state

Any additional thoughts or caveats?

Expected behavior

The above code works without error.

How to Reproduce

No response

Versions

Linux-5.15.85-penguins-rule-x86_64-with-glibc2.35
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0]
astropy 6.0.1
Numpy 1.24.2
pyerfa 2.0.1.4
Scipy 1.10.1
Matplotlib 3.7.1

[mhvk]

@Jerry-Ma - thanks for reporting! So, if I understand correctly, the actual problem is not specific to dill but really that an unbound info cannot be pickled, i.e., that the following fails:

import pickle
from astropy.utils.data_info import DataInfo
di = DataInfo()
pickle.dumps(di)

AttributeError: 'DataInfo' object has no attribute '_attrs'

I think your fix seems logical (though perhaps an alternative is to always set _attrs, maybe _attrs=None for the unbound case? your fix has the advantage of not changing any other behaviour, though). Would you be willing to make a PR? I think the test would just be with pickle on DataInfo rather than with dill on Masked - I don't think we want to add another dependency just for the tests if we can avoid it!

Aside - it is a bit curious that Masked instances can be pickled, i.e., pickle clearly does something a bit different from dill - this suggests that perhaps there is a bug in dill? Or is this an expected difference?