From bb28afef5c308010edda68091f435c77500c75ad Mon Sep 17 00:00:00 2001
From: Kaido Kert <kaidokert@gmail.com>
Date: Thu, 9 Jul 2015 19:16:51 -0700
Subject: [PATCH] Fix some of the #283 issues

---
 docs/wrapping.rst           |  3 +--
 pint/compat/__init__.py     |  5 +++++
 pint/testsuite/test_unit.py | 18 +++++++++---------
 pint/unit.py                |  9 ++++-----
 4 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/docs/wrapping.rst b/docs/wrapping.rst
index c46e1ea..b2ed162 100644
--- a/docs/wrapping.rst
+++ b/docs/wrapping.rst
@@ -150,7 +150,6 @@ To avoid the conversion of an argument or return value, use None
 Checking units
 ==============
 
-<to be written>
 When you want pint quantities to be used as inputs to your functions, pint provides a wrapper to ensure units are of
 correct type - or more precisely, they match the expected dimensionality of the physical quantity.
 
@@ -164,7 +163,7 @@ In the decorator format:
 
 .. doctest::
 
-    >>>@ureg.wraps('[length]')
+    >>>@ureg.check('[length]')
     ... def pendulum_period(length):
     ...     return 2*math.pi*math.sqrt(length/G)
 
diff --git a/pint/compat/__init__.py b/pint/compat/__init__.py
index 043f1bf..7442967 100644
--- a/pint/compat/__init__.py
+++ b/pint/compat/__init__.py
@@ -75,6 +75,11 @@ try:
 except ImportError:
     from .nullhandler import NullHandler
 
+try:
+    from itertools import zip_longest
+except ImportError:
+    from itertools import izip_longest as zip_longest
+
 try:
     import numpy as np
     from numpy import ndarray
diff --git a/pint/testsuite/test_unit.py b/pint/testsuite/test_unit.py
index 08ae9f9..edbb2fa 100644
--- a/pint/testsuite/test_unit.py
+++ b/pint/testsuite/test_unit.py
@@ -344,12 +344,12 @@ class TestRegistry(QuantityTestCase):
         f0 = ureg.check('[length]')(func)
         self.assertRaises(AttributeError, f0, 3.)
         self.assertEqual(f0(3. * ureg.centimeter), 0.03 * ureg.meter)
-        self.assertRaises(TypeError, f0, 3. * ureg.kilogram)
+        self.assertRaises(DimensionalityError, f0, 3. * ureg.kilogram)
 
         f0b = ureg.check(ureg.meter)(func)
         self.assertRaises(AttributeError, f0b, 3.)
         self.assertEqual(f0b(3. * ureg.centimeter), 0.03 * ureg.meter)
-        self.assertRaises(TypeError, f0b, 3. * ureg.kilogram)
+        self.assertRaises(DimensionalityError, f0b, 3. * ureg.kilogram)
 
         def gfunc(x, y):
             return x / y
@@ -360,19 +360,19 @@ class TestRegistry(QuantityTestCase):
 
         g1 = ureg.check('[speed]', '[time]')(gfunc)
         self.assertRaises(AttributeError, g1, 3.0, 1)
-        self.assertRaises(TypeError, g1, 1 * ureg.parsec, 1 * ureg.angstrom)
+        self.assertRaises(DimensionalityError, g1, 1 * ureg.parsec, 1 * ureg.angstrom)
         self.assertRaises(TypeError, g1, 1 * ureg.km / ureg.hour, 1 * ureg.hour, 3.0)
-        self.assertEquals(g1(3.6 * ureg.km / ureg.hour, 1 * ureg.second), 1 * ureg.meter / ureg.second ** 2)
+        self.assertEqual(g1(3.6 * ureg.km / ureg.hour, 1 * ureg.second), 1 * ureg.meter / ureg.second ** 2)
 
         g2 = ureg.check('[speed]')(gfunc)
         self.assertRaises(AttributeError, g2, 3.0, 1)
-        self.assertRaises(TypeError, g2, 2 * ureg.parsec)
-        self.assertRaises(TypeError, g2, 2 * ureg.parsec, 1.0)
-        self.assertEquals(g2(2.0 * ureg.km / ureg.hour, 2), 1 * ureg.km / ureg.hour)
+        self.assertRaises(DimensionalityError, g2, 2 * ureg.parsec)
+        self.assertRaises(DimensionalityError, g2, 2 * ureg.parsec, 1.0)
+        self.assertEqual(g2(2.0 * ureg.km / ureg.hour, 2), 1 * ureg.km / ureg.hour)
 
         g3 = ureg.check('[speed]', '[time]', '[mass]')(gfunc)
-        self.assertRaises(TypeError, g3, 1 * ureg.parsec, 1 * ureg.angstrom)
-        self.assertRaises(TypeError, g3, 1 * ureg.parsec, 1 * ureg.angstrom, 1 * ureg.kilogram)
+        self.assertRaises(DimensionalityError, g3, 1 * ureg.parsec, 1 * ureg.angstrom)
+        self.assertRaises(DimensionalityError, g3, 1 * ureg.parsec, 1 * ureg.angstrom, 1 * ureg.kilogram)
 
     def test_to_ref_vs_to(self):
         self.ureg.autoconvert_offset_to_baseunit = True
diff --git a/pint/unit.py b/pint/unit.py
index cbbc34e..7ca1591 100644
--- a/pint/unit.py
+++ b/pint/unit.py
@@ -29,7 +29,7 @@ from .util import (logger, pi_theorem, solve_dependencies, ParserHelper,
                    string_preprocessor, find_connected_nodes,
                    find_shortest_path, UnitsContainer, _is_dim,
                    SharedRegistryObject, to_units_container)
-from .compat import tokenizer, string_types, NUMERIC_TYPES, long_type
+from .compat import tokenizer, string_types, NUMERIC_TYPES, long_type, zip_longest
 from .definitions import (Definition, UnitDefinition, PrefixDefinition,
                           DimensionDefinition)
 from .converters import ScaleConverter
@@ -1166,11 +1166,10 @@ class UnitRegistry(object):
 
             @functools.wraps(func, assigned=assigned, updated=updated)
             def wrapper(*values, **kwargs):
-                for dim, value in itertools.izip_longest(dimensions, values):
+                for dim, value in zip_longest(dimensions, values):
                     if dim and value.dimensionality != dim:
-                        raise TypeError(
-                            'Expected units of %s, got %s' %
-                            (dim, value.dimensionality))
+                        raise DimensionalityError(value, 'a quantity of',
+                                                  value.dimensionality, dim)
                 return func(*values, **kwargs)
             return wrapper
         return decorator