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