1"""
2This file tests the behavior of subclasses of Representation and Frames
3"""
4from copy import deepcopy
5
6import pytest
7
8from astropy.coordinates import Longitude, Latitude
9from astropy.coordinates.representation import (REPRESENTATION_CLASSES,
10                                                SphericalRepresentation,
11                                                UnitSphericalRepresentation,
12                                                _invalidate_reprdiff_cls_hash)
13from astropy.coordinates.baseframe import frame_transform_graph
14from astropy.coordinates.transformations import FunctionTransform
15from astropy.coordinates import ICRS
16from astropy.coordinates.baseframe import RepresentationMapping
17
18import astropy.units as u
19
20import astropy.coordinates
21
22# Classes setup, borrowed from SunPy.
23
24# Here we define the classes *inside* the tests to make sure that we can wipe
25# the slate clean when the tests have finished running.
26
27
28def setup_function(func):
29    func.REPRESENTATION_CLASSES_ORIG = deepcopy(REPRESENTATION_CLASSES)
30
31
32def teardown_function(func):
33    REPRESENTATION_CLASSES.clear()
34    REPRESENTATION_CLASSES.update(func.REPRESENTATION_CLASSES_ORIG)
35    _invalidate_reprdiff_cls_hash()
36
37
38def test_unit_representation_subclass():
39
40    class Longitude180(Longitude):
41        def __new__(cls, angle, unit=None, wrap_angle=180*u.deg, **kwargs):
42            self = super().__new__(cls, angle, unit=unit, wrap_angle=wrap_angle,
43                                   **kwargs)
44            return self
45
46    class UnitSphericalWrap180Representation(UnitSphericalRepresentation):
47        attr_classes = {'lon': Longitude180,
48                        'lat': Latitude}
49
50    class SphericalWrap180Representation(SphericalRepresentation):
51        attr_classes = {'lon': Longitude180,
52                        'lat': Latitude,
53                        'distance': u.Quantity}
54
55        _unit_representation = UnitSphericalWrap180Representation
56
57    class MyFrame(ICRS):
58        default_representation = SphericalWrap180Representation
59        frame_specific_representation_info = {
60            'spherical': [
61                RepresentationMapping('lon', 'ra'),
62                RepresentationMapping('lat', 'dec')]
63        }
64        frame_specific_representation_info['unitsphericalwrap180'] = \
65            frame_specific_representation_info['sphericalwrap180'] = \
66            frame_specific_representation_info['spherical']
67
68    @frame_transform_graph.transform(FunctionTransform,
69                                     MyFrame, astropy.coordinates.ICRS)
70    def myframe_to_icrs(myframe_coo, icrs):
71        return icrs.realize_frame(myframe_coo._data)
72
73    f = MyFrame(10*u.deg, 10*u.deg)
74    assert isinstance(f._data, UnitSphericalWrap180Representation)
75    assert isinstance(f.ra, Longitude180)
76
77    g = f.transform_to(astropy.coordinates.ICRS())
78    assert isinstance(g, astropy.coordinates.ICRS)
79    assert isinstance(g._data, UnitSphericalWrap180Representation)
80
81    frame_transform_graph.remove_transform(MyFrame,
82                                           astropy.coordinates.ICRS,
83                                           None)
84