Skip to content

Commit fa0b684

Browse files
committed
Added preferred units support for .rescale
1 parent 1a0b450 commit fa0b684

2 files changed

Lines changed: 59 additions & 2 deletions

File tree

‎quantities/quantity.py‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@
1313
from .registry import unit_registry
1414
from .decorators import with_doc
1515

16+
PREFERRED = [] # List of preferred quanitities for each symbol,
17+
# e.g. PREFERRED = [pq.mV, pq.pA, pq.UnitQuantity('femtocoulomb', 1e-15*pq.C, 'fC')]
18+
# Intended to be overwritten in down-stream packages
19+
1620
if sys.version.startswith('3'):
1721
unicode = str
1822

@@ -195,10 +199,15 @@ def units(self, units):
195199
mag *= cf
196200
self._dimensionality = to_u.dimensionality
197201

198-
def rescale(self, units):
202+
def rescale(self, units=None):
199203
"""
200204
Return a copy of the quantity converted to the specified units
201205
"""
206+
if units is None:
207+
try:
208+
return self.rescale_preferred()
209+
except Exception as e:
210+
raise Exception('No argument passed to `.rescale` and %s' % e)
202211
to_dims = validate_dimensionality(units)
203212
if self.dimensionality == to_dims:
204213
return self.astype(self.dtype)
@@ -212,7 +221,15 @@ def rescale(self, units):
212221
%(from_u._dimensionality, to_u._dimensionality)
213222
)
214223
return Quantity(cf*self.magnitude, to_u)
215-
224+
225+
def rescale_preferred(self):
226+
units_str = str(self.simplified.dimensionality)
227+
for preferred in PREFERRED:
228+
if units_str == str(preferred.simplified.dimensionality):
229+
return self.rescale(preferred)
230+
raise Exception(("Preferred units for '%s' (or equivalent) not specified in "
231+
"quantites.quantity.PREFERRED." % self.dimensionality))
232+
216233
@with_doc(np.ndarray.astype)
217234
def astype(self, dtype=None, **kwargs):
218235
'''Scalars are returned as scalar Quantity arrays.'''

‎quantities/tests/test_conversion.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
# -*- coding: utf-8 -*-
22

3+
import unittest
34
from .. import units as pq
5+
from .. import quantity
46
from .common import TestCase
57

68

@@ -15,6 +17,44 @@ def test_inplace_conversion(self):
1517
def test_rescale(self):
1618
for u in ('ft', 'feet', pq.ft):
1719
self.assertQuantityEqual((10*pq.m).rescale(u), 32.80839895 * pq.ft)
20+
21+
def test_rescale_preferred(self):
22+
quantity.PREFERRED = [pq.mV, pq.pA]
23+
q = 10*pq.V
24+
self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.mV))
25+
q = 5*pq.A
26+
self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.pA))
27+
quantity.PREFERRED = []
28+
29+
def test_rescale_preferred_failure(self):
30+
quantity.PREFERRED = [pq.pA]
31+
q = 10*pq.V
32+
try:
33+
self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.mV))
34+
except:
35+
self.assertTrue(True)
36+
else:
37+
self.assertTrue(False)
38+
quantity.PREFERRED = []
39+
40+
def test_rescale_noargs(self):
41+
quantity.PREFERRED = [pq.mV, pq.pA]
42+
q = 10*pq.V
43+
self.assertQuantityEqual(q.rescale(), q.rescale(pq.mV))
44+
q = 5*pq.A
45+
self.assertQuantityEqual(q.rescale(), q.rescale(pq.pA))
46+
quantity.PREFERRED = []
47+
48+
def test_rescale_noargs_failure(self):
49+
quantity.PREFERRED = [pq.pA]
50+
q = 10*pq.V
51+
try:
52+
self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.mV))
53+
except:
54+
self.assertTrue(True)
55+
else:
56+
self.assertTrue(False)
57+
quantity.PREFERRED = []
1858

1959
def test_compound_reduction(self):
2060
pc_per_cc = pq.CompoundUnit("pc/cm**3")

0 commit comments

Comments
 (0)