Repository: astropy/astropy

Issue:
# Fix `PicklingError` when pickling `Masked` subclasses with initialized `info` attribute

<!-- These comments are hidden when you submit the pull request,
so you do not need to remove them! -->

<!-- Please be sure to check out our contributing guidelines,
https://github.com/astropy/astropy/blob/main/CONTRIBUTING.md .
Please be sure to check out our code of conduct,
https://github.com/astropy/astropy/blob/main/CODE_OF_CONDUCT.md . -->

<!-- If you are new or need to be re-acquainted with Astropy
contributing workflow, please see
https://docs.astropy.org/en/latest/development/quickstart.html .
There is even a practical example at
https://docs.astropy.org/en/latest/development/git_edit_workflow_examples.html . -->

<!-- Please just have a quick search on GitHub to see if a similar
pull request has already been posted.
We have old closed pull requests that might provide useful code or ideas
that directly tie in with your pull request. -->

<!-- We have several automatic features that run when a pull request is open.
They can appear daunting but do not worry because maintainers will help
you navigate them, if necessary. -->

### Description
<!-- Provide a general description of what your pull request does.
Complete the following sentence and add relevant details as you see fit. -->

<!-- In addition please ensure that the pull request title is descriptive
and allows maintainers to infer the applicable subpackage(s). -->

<!-- READ THIS FOR MANUAL BACKPORT FROM A MAINTAINER:
Apply "skip-basebranch-check" label **before** you open the PR! -->

This pull request is to address a `PicklingError` when attempting to pickle dynamically created `Masked` subclasses that have an initialized `info` attribute.  While #17685 ensured the `Masked` subclasses themselves were importable, their associated info classes, which are also created dynamically, remained inaccessible to `pickle`. This fix ensures that the info classes are assigned the correct module and are accessible through the module's `__getattr__` method.  This also resolves failures with pickling Masked `QTable` instances, since `MaskedQuantity` columns contain an `info` attribute. I also included regression tests for importing the info classes of common `Masked` subclasses.

Example of fixed behavior:
```python
import pickle
from astropy.utils.masked import Masked
from astropy import units as u

mq = Masked([1,2,3]*u.m, mask=[True,False,False])
mq.info # Access info to create it
# Previously raised _pickle.PicklingError
pickle.loads(pickle.dumps(mq))
```
<!-- If the pull request closes any open issues you can add this.
If you replace <Issue Number> with a number, GitHub will automatically link it.
If this pull request is unrelated to any issues, please remove
the following line. -->

Fixes #19040
Fixes #18206

<!-- Optional opt-out -->

- [ ] By checking this box, the PR author has requested that maintainers do **NOT** use the "Squash and Merge" button. Maintainers should respect this when possible; however, the final decision is at the discretion of the maintainer that merges the PR.


Implementation patch:
diff --git a/astropy/utils/masked/core.py b/astropy/utils/masked/core.py
index b87ac39ce429..4ad48e26fc8c 100644
--- a/astropy/utils/masked/core.py
+++ b/astropy/utils/masked/core.py
@@ -618,10 +618,13 @@ def __new__(newcls, *args, mask=None, **kwargs):
         if "info" not in cls.__dict__ and hasattr(cls._data_cls, "info"):
             data_info = cls._data_cls.info
             attr_names = data_info.attr_names | {"serialize_method"}
+            # Ensure the new info class uses the class's module.
+            # Without this, the DataInfoMeta metaclass incorrectly sets
+            # __module__ to its own.
             new_info = type(
                 cls.__name__ + "Info",
                 (MaskedArraySubclassInfo, data_info.__class__),
-                dict(attr_names=attr_names),
+                dict(attr_names=attr_names, __module__=cls.__module__),
             )
             cls.info = new_info()
 
@@ -1430,13 +1433,14 @@ def __repr__(self):
 
 
 def __getattr__(key):
-    """Make commonly used Masked subclasses importable for ASDF support.
+    """Make commonly used Masked subclasses and their info classes importable for ASDF support.
 
     Registered types associated with ASDF converters must be importable by
     their fully qualified name. Masked classes are dynamically created and have
     apparent names like ``astropy.utils.masked.core.MaskedQuantity`` although
     they aren't actually attributes of this module. Customize module attribute
-    lookup so that certain commonly used Masked classes are importable.
+    lookup so that certain commonly used Masked classes are importable. The same
+    is done for their dynamically created info classes.
 
     See:
     - https://asdf.readthedocs.io/en/latest/asdf/extending/converters.html#entry-point-performance-considerations
@@ -1450,13 +1454,14 @@ def __getattr__(key):
         base_class_name = key[len(Masked.__name__) :]
         for base_class_qualname in __construct_mixin_classes:
             module, _, name = base_class_qualname.rpartition(".")
-            if name == base_class_name:
+            is_info = name + "Info" == base_class_name
+            if name == base_class_name or is_info:
                 base_class = getattr(importlib.import_module(module), name)
                 # Try creating the masked class
                 masked_class = Masked(base_class)
                 # But only return it if it is a standard one, not one
                 # where we just used the ndarray fallback.
                 if base_class in Masked._masked_classes:
-                    return masked_class
+                    return masked_class.info.__class__ if is_info else masked_class
 
     raise AttributeError(f"module '{__name__}' has no attribute '{key}'")


Test patch:
diff --git a/astropy/table/tests/test_pickle.py b/astropy/table/tests/test_pickle.py
index 6e31e6de3f8e..8c670ff66101 100644
--- a/astropy/table/tests/test_pickle.py
+++ b/astropy/table/tests/test_pickle.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from astropy.coordinates import SkyCoord
+from astropy.coordinates import Angle, SkyCoord
 from astropy.table import Column, MaskedColumn, QTable, Table
 from astropy.table.table_helpers import simple_table
 from astropy.time import Time
@@ -141,6 +141,29 @@ def test_pickle_masked_table(protocol):
     assert tp.meta == t.meta
 
 
+def test_pickle_masked_qtable(protocol):
+    t = QTable(meta={"a": 1}, masked=True)
+    t["a"] = Quantity([1, 2], unit="m")
+    t["b"] = Angle([1, 2], unit=deg)
+
+    t["a"].mask[0] = True
+    t["b"].mask[1] = True
+
+    ts = pickle.dumps(t, protocol=protocol)
+    tp = pickle.loads(ts)
+
+    assert tp.__class__ is QTable
+
+    assert np.all(tp["a"].mask == t["a"].mask)
+    assert np.all(tp["a"].unmasked == t["a"].unmasked)
+    assert np.all(tp["b"].mask == t["b"].mask)
+    assert np.all(tp["b"].unmasked == t["b"].unmasked)
+    assert type(tp["a"]) is type(t["a"])  # nopep8
+    assert type(tp["b"]) is type(t["b"])  # nopep8
+    assert tp.meta == t.meta
+    assert type(tp) is type(t)
+
+
 def test_pickle_indexed_table(protocol):
     """
     Ensure that any indices that have been added will survive pickling.
diff --git a/astropy/utils/masked/tests/test_dynamic_subclasses.py b/astropy/utils/masked/tests/test_dynamic_subclasses.py
index f05a632ec9c3..972b09ac9177 100644
--- a/astropy/utils/masked/tests/test_dynamic_subclasses.py
+++ b/astropy/utils/masked/tests/test_dynamic_subclasses.py
@@ -1,5 +1,5 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
-"""Test importability of common Masked classes."""
+"""Test importability of common Masked classes and their info classes."""
 
 import pytest
 
@@ -11,3 +11,11 @@
 @pytest.mark.parametrize("base_class", [Quantity, Angle, Latitude, Longitude])
 def test_importable(base_class):
     assert getattr(core, f"Masked{base_class.__name__}") is core.Masked(base_class)
+
+
+@pytest.mark.parametrize("base_class", [Quantity, Angle, Latitude, Longitude])
+def test_importable_info(base_class):
+    assert (
+        getattr(core, f"Masked{base_class.__name__}Info")
+        is core.Masked(base_class).info.__class__
+    )
diff --git a/astropy/utils/masked/tests/test_pickle.py b/astropy/utils/masked/tests/test_pickle.py
new file mode 100644
index 000000000000..1429a91fbd43
--- /dev/null
+++ b/astropy/utils/masked/tests/test_pickle.py
@@ -0,0 +1,37 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+import pickle
+
+import numpy as np
+import pytest
+
+import astropy.units as u
+from astropy.coordinates import Angle, Latitude, Longitude
+from astropy.units import Quantity
+from astropy.utils.masked import Masked
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        Quantity([1, 2, 3], u.m),
+        Angle([1, 2, 3], u.deg),
+        Latitude([1, 2, 3], u.deg),
+        Longitude([1, 2, 3], u.deg),
+    ],
+)
+def test_masked_pickle(data):
+    mask = [True, False, False]
+    m = Masked(data, mask=mask)
+
+    # Force creation of the info object (see #19142)
+    assert m.info is not None
+
+    m2 = pickle.loads(pickle.dumps(m))
+
+    np.testing.assert_array_equal(m.unmasked, m2.unmasked)
+    np.testing.assert_array_equal(m.mask, m2.mask)
+
+    assert m.unit == m2.unit
+    assert type(m) is type(m2)
+    assert type(m.info) is type(m2.info)