1# -*- coding: utf-8 -*-
2
3import sys
4import pickle
5import copy
6
7from .. import units as pq
8from ..quantity import Quantity
9from ..uncertainquantity import UncertainQuantity
10from .. import constants
11from .common import TestCase
12
13
14class TestPersistence(TestCase):
15
16    def test_unitquantity_persistence(self):
17        x = pq.m
18        y = pickle.loads(pickle.dumps(x))
19        self.assertQuantityEqual(x, y)
20
21        x = pq.CompoundUnit("pc/cm**3")
22        y = pickle.loads(pickle.dumps(x))
23        self.assertQuantityEqual(x, y)
24
25    def test_quantity_persistence(self):
26        x = 20*pq.m
27        y = pickle.loads(pickle.dumps(x))
28        self.assertQuantityEqual(x, y)
29
30    def test_uncertainquantity_persistence(self):
31        x = UncertainQuantity(20, 'm', 0.2)
32        y = pickle.loads(pickle.dumps(x))
33        self.assertQuantityEqual(x, y)
34
35    def test_unitconstant_persistence(self):
36        x = constants.m_e
37        y = pickle.loads(pickle.dumps(x))
38        self.assertQuantityEqual(x, y)
39
40    def test_quantity_object_dtype(self):
41        # Regression test for github issue #113
42        x = Quantity(1,dtype=object)
43        y = pickle.loads(pickle.dumps(x))
44        self.assertQuantityEqual(x, y)
45
46    def test_uncertainquantity_object_dtype(self):
47        # Regression test for github issue #113
48        x = UncertainQuantity(20, 'm', 0.2, dtype=object)
49        y = pickle.loads(pickle.dumps(x))
50        self.assertQuantityEqual(x, y)
51
52    def test_backward_compat(self):
53        """ A few pickles collected before fixing #113 just to make sure we remain backwards compatible. """
54        orig = [
55            pq.m,
56            20*pq.m,
57            UncertainQuantity(20, 'm', 0.2),
58            constants.m_e,
59        ]
60        if sys.version_info[0] < 3:
61            data = [
62                # generated in python 2 with protocol=-1 (i.e. newest == 2)
63                b'\x80\x02cquantities.unitquantity\nUnitLength\nq\x00(U\x05meterq\x01NU\x01mq\x02N]q\x03(U\x06metersq\x04U\x05metreq\x05U\x06metresq\x06eNtq\x07Rq\x08K\x01K\x02K\x02\x86q\t\x86q\nb.',
64                b'\x80\x02cquantities.quantity\n_reconstruct_quantity\nq\x00(cquantities.quantity\nQuantity\nq\x01cnumpy\nndarray\nq\x02K\x00\x85q\x03U\x01bq\x04tq\x05Rq\x06(K\x01)cnumpy\ndtype\nq\x07U\x02f8q\x08K\x00K\x01\x87q\tRq\n(K\x03U\x01<q\x0bNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\x0cb\x89U\x08\x00\x00\x00\x00\x00\x004@q\rcquantities.dimensionality\nDimensionality\nq\x0e)\x81q\x0fcquantities.unitquantity\nUnitLength\nq\x10(U\x05meterq\x11NU\x01mq\x12N]q\x13(U\x06metersq\x14U\x05metreq\x15U\x06metresq\x16eNtq\x17Rq\x18K\x01K\x02K\x02\x86q\x19\x86q\x1abK\x01s}q\x1bbtq\x1cb.',
65                b'\x80\x02cquantities.quantity\n_reconstruct_quantity\nq\x00(cquantities.uncertainquantity\nUncertainQuantity\nq\x01cnumpy\nndarray\nq\x02K\x00\x85q\x03U\x01bq\x04tq\x05Rq\x06(K\x01)cnumpy\ndtype\nq\x07U\x02f8q\x08K\x00K\x01\x87q\tRq\n(K\x03U\x01<q\x0bNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\x0cb\x89U\x08\x00\x00\x00\x00\x00\x004@q\rcquantities.dimensionality\nDimensionality\nq\x0e)\x81q\x0fcquantities.unitquantity\nUnitLength\nq\x10(U\x05meterq\x11NU\x01mq\x12N]q\x13(U\x06metersq\x14U\x05metreq\x15U\x06metresq\x16eNtq\x17Rq\x18K\x01K\x02K\x02\x86q\x19\x86q\x1abK\x01s}q\x1bbh\x00(cquantities.quantity\nQuantity\nq\x1ch\x02h\x03h\x04tq\x1dRq\x1e(K\x01)h\n\x89U\x08\x9a\x99\x99\x99\x99\x99\xc9?q\x1fh\x0e)\x81q h\x18K\x01s}q!btq"btq#b.',
66                b'\x80\x02cquantities.unitquantity\nUnitConstant\nq\x00(U\relectron_massq\x01cquantities.quantity\n_reconstruct_quantity\nq\x02(cquantities.quantity\nQuantity\nq\x03cnumpy\nndarray\nq\x04K\x00\x85q\x05U\x01bq\x06tq\x07Rq\x08(K\x01)cnumpy\ndtype\nq\tU\x02f8q\nK\x00K\x01\x87q\x0bRq\x0c(K\x03U\x01<q\rNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\x0eb\x89U\x08N?\xeb\x93\xdcy\xb29q\x0fcquantities.dimensionality\nDimensionality\nq\x10)\x81q\x11cquantities.unitquantity\nUnitMass\nq\x12(U\x08kilogramq\x13NU\x02kgq\x14N]q\x15U\tkilogramsq\x16aNtq\x17Rq\x18K\x01K\x01K\x1b\x86q\x19\x86q\x1abK\x01s}q\x1bbtq\x1cbU\x03m_eq\x1dU\x04m\xe2\x82\x91q\x1e]q\x1fNtq Rq!K\x01K\x00M\x0f\x01\x86q"\x86q#b.',
67            ]
68        else:
69            data = [
70                # generated in python 3 with protocol=2
71                b'\x80\x02cquantities.unitquantity\nUnitLength\nq\x00(X\x05\x00\x00\x00meterq\x01NX\x01\x00\x00\x00mq\x02N]q\x03(X\x06\x00\x00\x00metersq\x04X\x05\x00\x00\x00metreq\x05X\x06\x00\x00\x00metresq\x06eNtq\x07Rq\x08K\x01K\x02K\x02\x86q\t\x86q\nb.',
72                b'\x80\x02cquantities.quantity\n_reconstruct_quantity\nq\x00(cquantities.quantity\nQuantity\nq\x01cnumpy\nndarray\nq\x02K\x00\x85q\x03X\x01\x00\x00\x00bq\x04tq\x05Rq\x06(K\x01)cnumpy\ndtype\nq\x07X\x02\x00\x00\x00f8q\x08K\x00K\x01\x87q\tRq\n(K\x03X\x01\x00\x00\x00<q\x0bNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\x0cb\x89c_codecs\nencode\nq\rX\x08\x00\x00\x00\x00\x00\x00\x00\x00\x004@q\x0eX\x06\x00\x00\x00latin1q\x0f\x86q\x10Rq\x11cquantities.dimensionality\nDimensionality\nq\x12)\x81q\x13cquantities.unitquantity\nUnitLength\nq\x14(X\x05\x00\x00\x00meterq\x15NX\x01\x00\x00\x00mq\x16N]q\x17(X\x06\x00\x00\x00metersq\x18X\x05\x00\x00\x00metreq\x19X\x06\x00\x00\x00metresq\x1aeNtq\x1bRq\x1cK\x01K\x02K\x02\x86q\x1d\x86q\x1ebK\x01stq\x1fb.',
73                b'\x80\x02cquantities.quantity\n_reconstruct_quantity\nq\x00(cquantities.uncertainquantity\nUncertainQuantity\nq\x01cnumpy\nndarray\nq\x02K\x00\x85q\x03X\x01\x00\x00\x00bq\x04tq\x05Rq\x06(K\x01)cnumpy\ndtype\nq\x07X\x02\x00\x00\x00f8q\x08K\x00K\x01\x87q\tRq\n(K\x03X\x01\x00\x00\x00<q\x0bNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\x0cb\x89c_codecs\nencode\nq\rX\x08\x00\x00\x00\x00\x00\x00\x00\x00\x004@q\x0eX\x06\x00\x00\x00latin1q\x0f\x86q\x10Rq\x11cquantities.dimensionality\nDimensionality\nq\x12)\x81q\x13cquantities.unitquantity\nUnitLength\nq\x14(X\x05\x00\x00\x00meterq\x15NX\x01\x00\x00\x00mq\x16N]q\x17(X\x06\x00\x00\x00metersq\x18X\x05\x00\x00\x00metreq\x19X\x06\x00\x00\x00metresq\x1aeNtq\x1bRq\x1cK\x01K\x02K\x02\x86q\x1d\x86q\x1ebK\x01sh\x00(cquantities.quantity\nQuantity\nq\x1fh\x02h\x03h\x04tq Rq!(K\x01)h\n\x89h\rX\x0f\x00\x00\x00\xc2\x9a\xc2\x99\xc2\x99\xc2\x99\xc2\x99\xc2\x99\xc3\x89?q"h\x0f\x86q#Rq$h\x12)\x81q%h\x1cK\x01stq&btq\'b.',
74                b'\x80\x02cquantities.unitquantity\nUnitConstant\nq\x00(X\r\x00\x00\x00electron_massq\x01cquantities.quantity\n_reconstruct_quantity\nq\x02(cquantities.quantity\nQuantity\nq\x03cnumpy\nndarray\nq\x04K\x00\x85q\x05X\x01\x00\x00\x00bq\x06tq\x07Rq\x08(K\x01)cnumpy\ndtype\nq\tX\x02\x00\x00\x00f8q\nK\x00K\x01\x87q\x0bRq\x0c(K\x03X\x01\x00\x00\x00<q\rNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\x0eb\x89c_codecs\nencode\nq\x0fX\x0c\x00\x00\x00N?\xc3\xab\xc2\x93\xc3\x9cy\xc2\xb29q\x10X\x06\x00\x00\x00latin1q\x11\x86q\x12Rq\x13cquantities.dimensionality\nDimensionality\nq\x14)\x81q\x15cquantities.unitquantity\nUnitMass\nq\x16(X\x08\x00\x00\x00kilogramq\x17NX\x02\x00\x00\x00kgq\x18N]q\x19X\t\x00\x00\x00kilogramsq\x1aaNtq\x1bRq\x1cK\x01K\x01K?\x86q\x1d\x86q\x1ebK\x01stq\x1fbX\x03\x00\x00\x00m_eq X\x04\x00\x00\x00m\xe2\x82\x91q!]q"Ntq#Rq$K\x01K\x00M!\x01\x86q%\x86q&b.',
75            ]
76        for x,d in zip(orig,data):
77            y = pickle.loads(d)
78            self.assertQuantityEqual(x, y)
79
80    def test_copy_quantity(self):
81        for dtype in [float,object]:
82            x = (20*pq.m).astype(dtype)
83            y = copy.copy(x)
84            self.assertQuantityEqual(x, y)
85
86    def test_copy_uncertainquantity(self):
87        for dtype in [float, object]:
88            x = UncertainQuantity(20, 'm', 0.2).astype(dtype)
89            y = copy.copy(x)
90            self.assertQuantityEqual(x, y)