1# Licensed under a 3-clause BSD style license - see LICENSE.rst 2 3""" 4This module contains functions for serializing core astropy objects via the 5YAML protocol. 6It provides functions `~astropy.io.misc.yaml.dump`, 7`~astropy.io.misc.yaml.load`, and `~astropy.io.misc.yaml.load_all` which 8call the corresponding functions in `PyYaml <https://pyyaml.org>`_ but use the 9`~astropy.io.misc.yaml.AstropyDumper` and `~astropy.io.misc.yaml.AstropyLoader` 10classes to define custom YAML tags for the following astropy classes: 11- `astropy.units.Unit` 12- `astropy.units.Quantity` 13- `astropy.time.Time` 14- `astropy.time.TimeDelta` 15- `astropy.coordinates.SkyCoord` 16- `astropy.coordinates.Angle` 17- `astropy.coordinates.Latitude` 18- `astropy.coordinates.Longitude` 19- `astropy.coordinates.EarthLocation` 20- `astropy.table.SerializedColumn` 21 22Example 23======= 24:: 25 >>> from astropy.io.misc import yaml 26 >>> import astropy.units as u 27 >>> from astropy.time import Time 28 >>> from astropy.coordinates import EarthLocation 29 >>> t = Time(2457389.0, format='mjd', 30 ... location=EarthLocation(1000, 2000, 3000, unit=u.km)) 31 >>> td = yaml.dump(t) 32 >>> print(td) 33 !astropy.time.Time 34 format: mjd 35 in_subfmt: '*' 36 jd1: 4857390.0 37 jd2: -0.5 38 location: !astropy.coordinates.earth.EarthLocation 39 ellipsoid: WGS84 40 x: !astropy.units.Quantity 41 unit: &id001 !astropy.units.Unit {unit: km} 42 value: 1000.0 43 y: !astropy.units.Quantity 44 unit: *id001 45 value: 2000.0 46 z: !astropy.units.Quantity 47 unit: *id001 48 value: 3000.0 49 out_subfmt: '*' 50 precision: 3 51 scale: utc 52 >>> ty = yaml.load(td) 53 >>> ty 54 <Time object: scale='utc' format='mjd' value=2457389.0> 55 >>> ty.location # doctest: +FLOAT_CMP 56 <EarthLocation (1000., 2000., 3000.) km> 57""" 58 59import base64 60 61import numpy as np 62import yaml 63 64from astropy.time import Time, TimeDelta 65from astropy import units as u 66from astropy import coordinates as coords 67from astropy.table import SerializedColumn 68 69 70__all__ = ['AstropyLoader', 'AstropyDumper', 'load', 'load_all', 'dump'] 71 72 73def _unit_representer(dumper, obj): 74 out = {'unit': str(obj.to_string())} 75 return dumper.represent_mapping('!astropy.units.Unit', out) 76 77 78def _unit_constructor(loader, node): 79 map = loader.construct_mapping(node) 80 return u.Unit(map['unit'], parse_strict='warn') 81 82 83def _serialized_column_representer(dumper, obj): 84 out = dumper.represent_mapping('!astropy.table.SerializedColumn', obj) 85 return out 86 87 88def _serialized_column_constructor(loader, node): 89 map = loader.construct_mapping(node) 90 return SerializedColumn(map) 91 92 93def _time_representer(dumper, obj): 94 out = obj.info._represent_as_dict() 95 return dumper.represent_mapping('!astropy.time.Time', out) 96 97 98def _time_constructor(loader, node): 99 map = loader.construct_mapping(node) 100 out = Time.info._construct_from_dict(map) 101 return out 102 103 104def _timedelta_representer(dumper, obj): 105 out = obj.info._represent_as_dict() 106 return dumper.represent_mapping('!astropy.time.TimeDelta', out) 107 108 109def _timedelta_constructor(loader, node): 110 map = loader.construct_mapping(node) 111 out = TimeDelta.info._construct_from_dict(map) 112 return out 113 114 115def _ndarray_representer(dumper, obj): 116 if not (obj.flags['C_CONTIGUOUS'] or obj.flags['F_CONTIGUOUS']): 117 obj = np.ascontiguousarray(obj) 118 119 if np.isfortran(obj): 120 obj = obj.T 121 order = 'F' 122 else: 123 order = 'C' 124 125 data_b64 = base64.b64encode(obj.tobytes()) 126 127 out = dict(buffer=data_b64, 128 dtype=str(obj.dtype), 129 shape=obj.shape, 130 order=order) 131 132 return dumper.represent_mapping('!numpy.ndarray', out) 133 134 135def _ndarray_constructor(loader, node): 136 map = loader.construct_mapping(node) 137 map['buffer'] = base64.b64decode(map['buffer']) 138 return np.ndarray(**map) 139 140 141def _quantity_representer(tag): 142 def representer(dumper, obj): 143 out = obj.info._represent_as_dict() 144 return dumper.represent_mapping(tag, out) 145 return representer 146 147 148def _quantity_constructor(cls): 149 def constructor(loader, node): 150 map = loader.construct_mapping(node) 151 return cls.info._construct_from_dict(map) 152 return constructor 153 154 155def _skycoord_representer(dumper, obj): 156 map = obj.info._represent_as_dict() 157 out = dumper.represent_mapping('!astropy.coordinates.sky_coordinate.SkyCoord', 158 map) 159 return out 160 161 162def _skycoord_constructor(loader, node): 163 map = loader.construct_mapping(node) 164 out = coords.SkyCoord.info._construct_from_dict(map) 165 return out 166 167 168# Straight from yaml's Representer 169def _complex_representer(self, data): 170 if data.imag == 0.0: 171 data = f'{data.real!r}' 172 elif data.real == 0.0: 173 data = f'{data.imag!r}j' 174 elif data.imag > 0: 175 data = f'{data.real!r}+{data.imag!r}j' 176 else: 177 data = f'{data.real!r}{data.imag!r}j' 178 return self.represent_scalar('tag:yaml.org,2002:python/complex', data) 179 180 181def _complex_constructor(loader, node): 182 map = loader.construct_scalar(node) 183 return complex(map) 184 185 186class AstropyLoader(yaml.SafeLoader): 187 """ 188 Custom SafeLoader that constructs astropy core objects as well 189 as Python tuple and unicode objects. 190 191 This class is not directly instantiated by user code, but instead is 192 used to maintain the available constructor functions that are 193 called when parsing a YAML stream. See the `PyYaml documentation 194 <https://pyyaml.org/wiki/PyYAMLDocumentation>`_ for details of the 195 class signature. 196 """ 197 198 def _construct_python_tuple(self, node): 199 return tuple(self.construct_sequence(node)) 200 201 def _construct_python_unicode(self, node): 202 return self.construct_scalar(node) 203 204 205class AstropyDumper(yaml.SafeDumper): 206 """ 207 Custom SafeDumper that represents astropy core objects as well 208 as Python tuple and unicode objects. 209 210 This class is not directly instantiated by user code, but instead is 211 used to maintain the available representer functions that are 212 called when generating a YAML stream from an object. See the 213 `PyYaml documentation <https://pyyaml.org/wiki/PyYAMLDocumentation>`_ 214 for details of the class signature. 215 """ 216 217 def _represent_tuple(self, data): 218 return self.represent_sequence('tag:yaml.org,2002:python/tuple', data) 219 220 221AstropyDumper.add_multi_representer(u.UnitBase, _unit_representer) 222AstropyDumper.add_multi_representer(u.FunctionUnitBase, _unit_representer) 223AstropyDumper.add_representer(tuple, AstropyDumper._represent_tuple) 224AstropyDumper.add_representer(np.ndarray, _ndarray_representer) 225AstropyDumper.add_representer(Time, _time_representer) 226AstropyDumper.add_representer(TimeDelta, _timedelta_representer) 227AstropyDumper.add_representer(coords.SkyCoord, _skycoord_representer) 228AstropyDumper.add_representer(SerializedColumn, _serialized_column_representer) 229 230# Numpy dtypes 231AstropyDumper.add_representer(np.bool_, yaml.representer.SafeRepresenter.represent_bool) 232for np_type in [np.int_, np.intc, np.intp, np.int8, np.int16, np.int32, 233 np.int64, np.uint8, np.uint16, np.uint32, np.uint64]: 234 AstropyDumper.add_representer(np_type, 235 yaml.representer.SafeRepresenter.represent_int) 236for np_type in [np.float_, np.float16, np.float32, np.float64, 237 np.longdouble]: 238 AstropyDumper.add_representer(np_type, 239 yaml.representer.SafeRepresenter.represent_float) 240for np_type in [np.complex_, complex, np.complex64, np.complex128]: 241 AstropyDumper.add_representer(np_type, _complex_representer) 242 243AstropyLoader.add_constructor('tag:yaml.org,2002:python/complex', 244 _complex_constructor) 245AstropyLoader.add_constructor('tag:yaml.org,2002:python/tuple', 246 AstropyLoader._construct_python_tuple) 247AstropyLoader.add_constructor('tag:yaml.org,2002:python/unicode', 248 AstropyLoader._construct_python_unicode) 249AstropyLoader.add_constructor('!astropy.units.Unit', _unit_constructor) 250AstropyLoader.add_constructor('!numpy.ndarray', _ndarray_constructor) 251AstropyLoader.add_constructor('!astropy.time.Time', _time_constructor) 252AstropyLoader.add_constructor('!astropy.time.TimeDelta', _timedelta_constructor) 253AstropyLoader.add_constructor('!astropy.coordinates.sky_coordinate.SkyCoord', 254 _skycoord_constructor) 255AstropyLoader.add_constructor('!astropy.table.SerializedColumn', 256 _serialized_column_constructor) 257 258for cls, tag in ((u.Quantity, '!astropy.units.Quantity'), 259 (u.Magnitude, '!astropy.units.Magnitude'), 260 (u.Dex, '!astropy.units.Dex'), 261 (u.Decibel, '!astropy.units.Decibel'), 262 (coords.Angle, '!astropy.coordinates.Angle'), 263 (coords.Latitude, '!astropy.coordinates.Latitude'), 264 (coords.Longitude, '!astropy.coordinates.Longitude'), 265 (coords.EarthLocation, '!astropy.coordinates.earth.EarthLocation')): 266 AstropyDumper.add_multi_representer(cls, _quantity_representer(tag)) 267 AstropyLoader.add_constructor(tag, _quantity_constructor(cls)) 268 269for cls in (list(coords.representation.REPRESENTATION_CLASSES.values()) 270 + list(coords.representation.DIFFERENTIAL_CLASSES.values())): 271 name = cls.__name__ 272 # Add representations/differentials defined in astropy. 273 if name in coords.representation.__all__: 274 tag = '!astropy.coordinates.' + name 275 AstropyDumper.add_multi_representer(cls, _quantity_representer(tag)) 276 AstropyLoader.add_constructor(tag, _quantity_constructor(cls)) 277 278 279def load(stream): 280 """Parse the first YAML document in a stream using the AstropyLoader and 281 produce the corresponding Python object. 282 283 Parameters 284 ---------- 285 stream : str or file-like 286 YAML input 287 288 Returns 289 ------- 290 obj : object 291 Object corresponding to YAML document 292 """ 293 return yaml.load(stream, Loader=AstropyLoader) 294 295 296def load_all(stream): 297 """Parse the all YAML documents in a stream using the AstropyLoader class and 298 produce the corresponding Python object. 299 300 Parameters 301 ---------- 302 stream : str or file-like 303 YAML input 304 305 Returns 306 ------- 307 obj : object 308 Object corresponding to YAML document 309 310 """ 311 return yaml.load_all(stream, Loader=AstropyLoader) 312 313 314def dump(data, stream=None, **kwargs): 315 """Serialize a Python object into a YAML stream using the AstropyDumper class. 316 If stream is None, return the produced string instead. 317 318 Parameters 319 ---------- 320 data: object 321 Object to serialize to YAML 322 stream : file-like, optional 323 YAML output (if not supplied a string is returned) 324 **kwargs 325 Other keyword arguments that get passed to yaml.dump() 326 327 Returns 328 ------- 329 out : str or None 330 If no ``stream`` is supplied then YAML output is returned as str 331 332 """ 333 kwargs['Dumper'] = AstropyDumper 334 kwargs.setdefault('default_flow_style', None) 335 return yaml.dump(data, stream=stream, **kwargs) 336