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