Merge pull request #464 from tecki/feat/dimensionless-wraps

support dimensionless values in wraps
hgrecco · Nov 22, 2016 · de8fdff · de8fdff
2 parents 3a98ac2 + 9a42cd8
commit de8fdff
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/pint/registry_helpers.py b/pint/registry_helpers.py
@@ -13,7 +13,7 @@
 
 from .compat import string_types, zip_longest
 from .errors import DimensionalityError
-from .util import to_units_container
+from .util import to_units_container, UnitsContainer
 
 
 def _replace_units(original_units, values_by_name):
@@ -26,7 +26,7 @@ def _replace_units(original_units, values_by_name):
     for arg_name, exponent in original_units.items():
         q = q * values_by_name[arg_name] ** exponent
 
-    return to_units_container(q)
+    return getattr(q, "_units", UnitsContainer({}))
 
 
 def _to_units_container(a):
@@ -95,13 +95,16 @@ def _converter(ureg, values, strict):
 
         # first pass: Grab named values
         for ndx in defs_args_ndx:
-            values_by_name[args_as_uc[ndx][0]] = values[ndx]
-            new_values[ndx] = values[ndx]._magnitude
+            value = values[ndx]
+            values_by_name[args_as_uc[ndx][0]] = value
+            new_values[ndx] = getattr(value, "_magnitude", value)
 
         # second pass: calculate derived values based on named values
         for ndx in dependent_args_ndx:
-            new_values[ndx] = ureg._convert(values[ndx]._magnitude,
-                                            values[ndx]._units,
+            value = values[ndx]
+            assert _replace_units(args_as_uc[ndx][0], values_by_name) is not None
+            new_values[ndx] = ureg._convert(getattr(value, "_magnitude", value),
+                                            getattr(value, "_units", UnitsContainer({})),
                                             _replace_units(args_as_uc[ndx][0], values_by_name))
 
         # third pass: convert other arguments

diff --git a/pint/testsuite/test_unit.py b/pint/testsuite/test_unit.py
@@ -364,6 +364,7 @@ def gfunc3(x, y):
 
         g0 = ureg.wraps('=A', ['=A', '=A'])(gfunc)
         self.assertEqual(g0(3. * ureg.meter, 1. * ureg.centimeter), rst.to('meter'))
+        self.assertEqual(g0(3, 1), 4)
 
         g1 = ureg.wraps('=A', ['=A', '=A'])(gfunc)
         self.assertEqual(g1(3. * ureg.meter, 1. * ureg.centimeter), rst.to('centimeter'))
@@ -375,9 +376,12 @@ def gfunc3(x, y):
         a = 3. * ureg.meter
         b = (2. * ureg.centimeter) ** 2
         self.assertEqual(g3(a, b), gfunc2(a, b))
+        self.assertEqual(g3(3, 2), gfunc2(3, 2))
 
         g4 = ureg.wraps('=A**2 * B', ['=A', '=B'])(gfunc3)
         self.assertEqual(g4(3. * ureg.meter, 2. * ureg.second), ureg('(3*meter)**2 * 2 *second'))
+        self.assertEqual(g4(3. * ureg.meter, 2.), ureg('(3*meter)**2 * 2'))
+        self.assertEqual(g4(3., 2. * ureg.second), ureg('3**2 * 2 * second'))
 
 
     def test_check(self):