1"""
2"""
3from __future__ import absolute_import
4
5import sys
6import weakref
7
8import numpy
9
10from .dimensionality import Dimensionality
11from . import markup
12from .quantity import Quantity, get_conversion_factor
13from .registry import unit_registry
14from .decorators import memoize, with_doc
15
16
17__all__ = [
18    'CompoundUnit', 'Dimensionless', 'UnitConstant', 'UnitCurrency',
19    'UnitCurrent', 'UnitInformation', 'UnitLength', 'UnitLuminousIntensity',
20    'UnitMass', 'UnitMass', 'UnitQuantity', 'UnitSubstance', 'UnitTemperature',
21    'UnitTime', 'set_default_units'
22]
23
24
25class UnitQuantity(Quantity):
26
27    _primary_order = 90
28    _secondary_order = 0
29    _reference_quantity = None
30
31    __array_priority__ = 20
32
33    def __new__(
34        cls, name, definition=None, symbol=None, u_symbol=None,
35        aliases=[], doc=None
36    ):
37        try:
38            assert isinstance(name, str)
39        except AssertionError:
40            raise TypeError('name must be a string, got %s (not unicode)'%name)
41        try:
42            assert symbol is None or isinstance(symbol, str)
43        except AssertionError:
44            raise TypeError(
45                'symbol must be a string, '
46                'got %s (u_symbol can be unicode)'%symbol
47            )
48
49        ret = numpy.array(1, dtype='d').view(cls)
50        ret.flags.writeable = False
51
52        ret._name = name
53        ret._symbol = symbol
54        ret._u_symbol = u_symbol
55        if doc is not None:
56            ret.__doc__ = doc
57
58        if definition is not None:
59            if not isinstance(definition, Quantity):
60                definition *= dimensionless
61            ret._definition = definition
62            ret._conv_ref = definition._reference
63        else:
64            ret._definition = None
65            ret._conv_ref = None
66
67        ret._aliases = aliases
68
69        ret._format_order = (ret._primary_order, ret._secondary_order)
70        ret.__class__._secondary_order += 1
71
72        return ret
73
74    def __init__(
75        self, name, definition=None, symbol=None, u_symbol=None,
76        aliases=[], doc=None
77    ):
78        unit_registry[name] = self
79        if symbol:
80            unit_registry[symbol] = self
81        for alias in aliases:
82            unit_registry[alias] = self
83
84    def __array_finalize__(self, obj):
85        pass
86
87    def __hash__(self):
88        return hash((type(self), self._name))
89
90    @property
91    def _reference(self):
92        if self._conv_ref is None:
93            return self
94        else:
95            return self._conv_ref
96
97    @property
98    def _dimensionality(self):
99        return Dimensionality({self:1})
100
101    @property
102    def format_order(self):
103        return self._format_order
104
105    @property
106    def name(self):
107        return self._name
108
109    @property
110    def definition(self):
111        if self._definition is None:
112            return self
113        else:
114            return self._definition
115
116    @property
117    def simplified(self):
118        return self._reference.simplified
119
120    @property
121    def symbol(self):
122        if self._symbol:
123            return self._symbol
124        else:
125            return self.name
126
127    @property
128    def u_symbol(self):
129        if self._u_symbol:
130            return self._u_symbol
131        else:
132            return self.symbol
133
134    @property
135    def units(self):
136        return self
137    @units.setter
138    def units(self, units):
139        raise AttributeError('can not modify protected units')
140
141    def __repr__(self):
142        ref = self._definition
143        if ref:
144            ref = ', %s * %s'%(str(ref.magnitude), ref.dimensionality.string)
145        else:
146            ref = ''
147        symbol = self._symbol
148        symbol = ', %s'%(repr(symbol)) if symbol else ''
149        if markup.config.use_unicode:
150            u_symbol = self._u_symbol
151            u_symbol = ', %s'%(repr(u_symbol)) if u_symbol else ''
152        else:
153            u_symbol = ''
154        return '%s(%s%s%s%s)'%(
155            self.__class__.__name__, repr(self.name), ref, symbol, u_symbol
156        )
157
158    @with_doc(Quantity.__str__, use_header=False)
159    def __str__(self):
160        if self.u_symbol != self.name:
161            if markup.config.use_unicode:
162                s = '1 %s (%s)'%(self.u_symbol, self.name)
163            else:
164                s = '1 %s (%s)'%(self.symbol, self.name)
165        else:
166            s = '1 %s'%self.name
167
168        return s
169
170    @with_doc(Quantity.__add__, use_header=False)
171    def __add__(self, other):
172        return self.view(Quantity).__add__(other)
173
174    @with_doc(Quantity.__radd__, use_header=False)
175    def __radd__(self, other):
176        try:
177            return self.rescale(other.units).__radd__(other)
178        except AttributeError:
179            return self.view(Quantity).__radd__(other)
180
181    @with_doc(Quantity.__sub__, use_header=False)
182    def __sub__(self, other):
183        return self.view(Quantity).__sub__(other)
184
185    @with_doc(Quantity.__rsub__, use_header=False)
186    def __rsub__(self, other):
187        try:
188            return self.rescale(other.units).__rsub__(other)
189        except AttributeError:
190            return self.view(Quantity).__rsub__(other)
191
192    @with_doc(Quantity.__mod__, use_header=False)
193    def __mod__(self, other):
194        return self.view(Quantity).__mod__(other)
195
196    @with_doc(Quantity.__rsub__, use_header=False)
197    def __rmod__(self, other):
198        try:
199            return self.rescale(other.units).__rmod__(other)
200        except AttributeError:
201            return self.view(Quantity).__rmod__(other)
202
203    @with_doc(Quantity.__mul__, use_header=False)
204    def __mul__(self, other):
205        return self.view(Quantity).__mul__(other)
206
207    @with_doc(Quantity.__rmul__, use_header=False)
208    def __rmul__(self, other):
209        return self.view(Quantity).__rmul__(other)
210
211    @with_doc(Quantity.__truediv__, use_header=False)
212    def __truediv__(self, other):
213        return self.view(Quantity).__truediv__(other)
214
215    @with_doc(Quantity.__rtruediv__, use_header=False)
216    def __rtruediv__(self, other):
217        return self.view(Quantity).__rtruediv__(other)
218
219    if sys.version_info[0] < 3:
220        @with_doc(Quantity.__div__, use_header=False)
221        def __div__(self, other):
222            return self.view(Quantity).__div__(other)
223
224        @with_doc(Quantity.__rdiv__, use_header=False)
225        def __rdiv__(self, other):
226            return self.view(Quantity).__rdiv__(other)
227
228    @with_doc(Quantity.__pow__, use_header=False)
229    def __pow__(self, other):
230        return self.view(Quantity).__pow__(other)
231
232    @with_doc(Quantity.__rpow__, use_header=False)
233    def __rpow__(self, other):
234        return self.view(Quantity).__rpow__(other)
235
236    @with_doc(Quantity.__iadd__, use_header=False)
237    def __iadd__(self, other):
238        raise TypeError('can not modify protected units')
239
240    @with_doc(Quantity.__isub__, use_header=False)
241    def __isub__(self, other):
242        raise TypeError('can not modify protected units')
243
244    @with_doc(Quantity.__imul__, use_header=False)
245    def __imul__(self, other):
246        raise TypeError('can not modify protected units')
247
248    @with_doc(Quantity.__itruediv__, use_header=False)
249    def __itruediv__(self, other):
250        raise TypeError('can not modify protected units')
251
252    if sys.version_info[0] < 3:
253        @with_doc(Quantity.__idiv__, use_header=False)
254        def __idiv__(self, other):
255            raise TypeError('can not modify protected units')
256
257    @with_doc(Quantity.__ipow__, use_header=False)
258    def __ipow__(self, other):
259        raise TypeError('can not modify protected units')
260
261    def __getstate__(self):
262        """
263        Return the internal state of the quantity, for pickling
264        purposes.
265
266        """
267        state = (1, self._format_order)
268        return state
269
270    def __setstate__(self, state):
271        ver, fo = state
272        self._format_order = fo
273
274    def __reduce__(self):
275        """
276        Return a tuple for pickling a UnitQuantity.
277        """
278        return (
279            type(self),
280            (
281                self._name,
282                self._definition,
283                self._symbol,
284                self._u_symbol,
285                self._aliases,
286                self.__doc__
287            ),
288            self.__getstate__()
289        )
290
291    def copy(self):
292        return (
293            type(self)(
294                self._name,
295                self._definition,
296                self._symbol,
297                self._u_symbol,
298                self._aliases,
299                self.__doc__
300                )
301            )
302
303unit_registry['UnitQuantity'] = UnitQuantity
304
305
306class IrreducibleUnit(UnitQuantity):
307
308    _default_unit = None
309
310    def __init__(
311        self, name, definition=None, symbol=None, u_symbol=None,
312        aliases=[], doc=None
313    ):
314        super(IrreducibleUnit, self).__init__(
315            name, definition, symbol, u_symbol, aliases, doc
316        )
317        cls = type(self)
318        if cls._default_unit is None:
319            cls._default_unit = self
320
321    @property
322    def simplified(self):
323        return self.view(Quantity).rescale(self.get_default_unit())
324
325    @classmethod
326    def get_default_unit(cls):
327        return cls._default_unit
328    @classmethod
329    def set_default_unit(cls, unit):
330        if unit is None:
331            return
332        if isinstance(unit, str):
333            unit = unit_registry[unit]
334        try:
335            # check that conversions are possible:
336            get_conversion_factor(cls._default_unit, unit)
337        except ValueError:
338            raise TypeError('default unit must be of same type')
339        cls._default_unit = unit
340
341
342class UnitMass(IrreducibleUnit):
343
344    _primary_order = 1
345
346
347class UnitLength(IrreducibleUnit):
348
349    _primary_order = 2
350
351
352class UnitTime(IrreducibleUnit):
353
354    _primary_order = 3
355
356
357class UnitCurrent(IrreducibleUnit):
358
359    _primary_order = 4
360
361
362class UnitLuminousIntensity(IrreducibleUnit):
363
364    _primary_order = 5
365
366
367class UnitSubstance(IrreducibleUnit):
368
369    _primary_order = 6
370
371
372class UnitTemperature(IrreducibleUnit):
373
374    _primary_order = 7
375
376
377class UnitInformation(IrreducibleUnit):
378
379    _primary_order = 8
380
381
382class UnitCurrency(IrreducibleUnit):
383
384    _primary_order = 9
385
386
387class CompoundUnit(UnitQuantity):
388
389    _primary_order = 99
390
391    def __new__(cls, name):
392        return UnitQuantity.__new__(cls, name, unit_registry[name])
393
394    def __init__(self, name):
395        # do not register
396        return
397
398    @with_doc(UnitQuantity.__add__, use_header=False)
399    def __repr__(self):
400        return '1 %s'%self.name
401
402    @property
403    def name(self):
404        if markup.config.use_unicode:
405            return '(%s)'%(markup.superscript(self._name))
406        else:
407            return '(%s)'%self._name
408
409    def __reduce__(self):
410        """
411        Return a tuple for pickling a UnitQuantity.
412        """
413        return (
414            type(self),
415            (self._name, ),
416            self.__getstate__()
417            )
418
419    def copy(self):
420        return type(self)(self._name)
421
422unit_registry['CompoundUnit'] = CompoundUnit
423
424
425class Dimensionless(UnitQuantity):
426
427    _primary_order = 100
428
429    def __init__(self, name, definition=None):
430        self._name = name
431
432        if definition is None:
433            definition = self
434        self._definition = definition
435
436        self._format_order = (self._primary_order, self._secondary_order)
437        self.__class__._secondary_order += 1
438
439        unit_registry[name] = self
440
441    def __reduce__(self):
442        """
443        Return a tuple for pickling a UnitQuantity.
444        """
445        return (
446            type(self),
447            (
448                self._name,
449            ),
450            self.__getstate__()
451        )
452
453    @property
454    def _dimensionality(self):
455        return Dimensionality()
456
457dimensionless = Dimensionless('dimensionless')
458
459
460class UnitConstant(UnitQuantity):
461
462    _primary_order = 0
463
464    def __init__(
465        self, name, definition=None, symbol=None, u_symbol=None,
466        aliases=[], doc=None
467    ):
468        # we dont want to register constants in the unit registry
469        return
470
471
472def set_default_units(
473    system=None, currency=None, current=None, information=None, length=None,
474    luminous_intensity=None, mass=None, substance=None, temperature=None,
475    time=None
476):
477    """
478    Set the default units in which simplified quantities will be
479    expressed.
480
481    system sets the unit system, and can be "SI" or "cgs". All other
482    keyword arguments will accept either a string or a unit quantity.
483    An error will be raised if it is not possible to convert between
484    old and new defaults, so it is not possible to set "kg" as the
485    default unit for time.
486
487    If both system and individual defaults are given, the system
488    defaults will be applied first, followed by the individual ones.
489    """
490    if system is not None:
491        system = system.lower()
492        try:
493            assert system in ('si', 'cgs')
494        except AssertionError:
495            raise ValueError('system must be "SI" or "cgs", got "%s"' % system)
496        if system == 'si':
497            UnitCurrent.set_default_unit('A')
498            UnitLength.set_default_unit('m')
499            UnitMass.set_default_unit('kg')
500        elif system == 'cgs':
501            UnitLength.set_default_unit('cm')
502            UnitMass.set_default_unit('g')
503        UnitLuminousIntensity.set_default_unit('cd')
504        UnitSubstance.set_default_unit('mol')
505        UnitTemperature.set_default_unit('degK')
506        UnitTime.set_default_unit('s')
507
508    UnitCurrency.set_default_unit(currency)
509    UnitCurrent.set_default_unit(current)
510    UnitInformation.set_default_unit(information)
511    UnitLength.set_default_unit(length)
512    UnitLuminousIntensity.set_default_unit(luminous_intensity)
513    UnitMass.set_default_unit(mass)
514    UnitSubstance.set_default_unit(substance)
515    UnitTemperature.set_default_unit(temperature)
516    UnitTime.set_default_unit(time)
517