Merge pull request #5021 from eteq/fix-frame-attr-set-delayed

astrofrog · eteq · commit ba428ef66d0e · 2016-06-06T18:27:34.000-04:00
Fix SkyCoord AttributeError when a new frame is created after a SkyCoord has been inited
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -1538,6 +1538,10 @@ Bug Fixes
 
 - ``astropy.coordinates``
 
+  - ``SkyCoord`` objects created before a new frame which has frame attributes
+    is created no longer raise ``AttributeError`` when the new attributes are
+    accessed [#5021]
+
   - Fix some errors in the implementation of aberration  for ``get_sun``. [#4979]
 
 - ``astropy.cosmology``
diff --git a/astropy/coordinates/sky_coordinate.py b/astropy/coordinates/sky_coordinate.py
@@ -424,7 +424,17 @@ def __getattr__(self, attr):
                 if attr in self.frame.get_frame_attr_names():
                     return getattr(self.frame, attr)
                 else:
-                    return getattr(self, '_' + attr)
+                    try:
+                        return getattr(self, '_' + attr)
+                    except AttributeError:
+                        # this can happen because FRAME_ATTR_NAMES_SET is
+                        # dynamic.  So if a frame is added to the transform
+                        # graph after this SkyCoord was created, the "real"
+                        # underlying attribute - e.g. `_equinox` does not exist
+                        # on the SkyCoord.  So we add this case to just use
+                        # None, as it wouldn't have been possible for the user
+                        # to have set the value until the frame existed anyway
+                        return None
 
             # Some attributes might not fall in the above category but still
             # are available through self._sky_coord_frame.
diff --git a/astropy/coordinates/tests/test_sky_coord.py b/astropy/coordinates/tests/test_sky_coord.py
@@ -21,7 +21,9 @@
 from ..representation import REPRESENTATION_CLASSES
 from ...coordinates import (ICRS, FK4, FK5, Galactic, SkyCoord, Angle,
                             SphericalRepresentation, CartesianRepresentation,
-                            UnitSphericalRepresentation, AltAz)
+                            UnitSphericalRepresentation, AltAz,
+                            BaseCoordinateFrame, FrameAttribute,
+                            frame_transform_graph)
 from ...coordinates import Latitude, EarthLocation
 from ...time import Time
 from ...utils import minversion
@@ -1214,3 +1216,39 @@ def test_spherical_offsets():
     dra, ddec = i00.spherical_offsets_to(i1deg)
     assert_allclose(dra, 1*u.deg)
     assert_allclose(ddec, 1*u.deg)
+
+def test_frame_attr_changes():
+    """
+    This tests the case where a frame is added with a new frame attribute after
+    a SkyCoord has been created.  This is necessary because SkyCoords get the
+    attributes set at creation time, but the set of attributes can change as
+    frames are added or removed from the transform graph.  This makes sure that
+    everything continues to work consistently.
+    """
+    sc_before = SkyCoord(1*u.deg, 2*u.deg, frame='icrs')
+
+    assert 'fakeattr' not in dir(sc_before)
+
+    class FakeFrame(BaseCoordinateFrame):
+        fakeattr = FrameAttribute()
+
+    # doesn't matter what this does as long as it just puts the frame in the
+    # transform graph
+    transset = (ICRS, FakeFrame, lambda c,f:c)
+    frame_transform_graph.add_transform(*transset)
+    try:
+        assert 'fakeattr' in dir(sc_before)
+        assert sc_before.fakeattr is None
+
+        sc_after1 = SkyCoord(1*u.deg, 2*u.deg, frame='icrs')
+        assert 'fakeattr' in dir(sc_after1)
+        assert sc_after1.fakeattr is None
+
+        sc_after2 = SkyCoord(1*u.deg, 2*u.deg, frame='icrs', fakeattr=1)
+        assert sc_after2.fakeattr == 1
+    finally:
+        frame_transform_graph.remove_transform(*transset)
+
+    assert 'fakeattr' not in dir(sc_before)
+    assert 'fakeattr' not in dir(sc_after1)
+    assert 'fakeattr' not in dir(sc_after2)
