Repository: ModellingWebLab/cellmlmanip

Issue:
Add method to obtain unit name
Given a `pint.unit` we'd like to know its name. This used to be possible with `str(unit)` but that now adds a prefix we need to strip

Discussion:

Patch:
diff --git a/cellmlmanip/model.py b/cellmlmanip/model.py
index de9aec1..fae9c45 100644
--- a/cellmlmanip/model.py
+++ b/cellmlmanip/model.py
@@ -278,8 +278,8 @@ class Model(object):
         lhs_units = self.units.evaluate_units(equality.lhs)
         rhs_units = self.units.evaluate_units(equality.rhs)
         assert self.units.is_equivalent(rhs_units, lhs_units), 'Units %s %s != %s %s' % (
-            lhs_units, self.units.show_base_units(lhs_units),
-            rhs_units, self.units.show_base_units(rhs_units)
+            self.units.format(lhs_units), self.units.format(lhs_units, True),
+            self.units.format(rhs_units), self.units.format(rhs_units, True),
         )
 
     def get_equations_for(self, symbols, recurse=True, strip_units=True):
diff --git a/cellmlmanip/units.py b/cellmlmanip/units.py
index 691a5ed..a6ea6c0 100644
--- a/cellmlmanip/units.py
+++ b/cellmlmanip/units.py
@@ -57,6 +57,9 @@ _TRIG_FUNCTIONS = {
 # It says (1) not preceded by a number or period, (2) then at least one ascii letter, (3) then numbers or _ allowed too
 _WORD = re.compile('(?<![0-9.])[a-zA-Z_]+[a-zA-Z0-9_]*')
 
+# Regex to remove UnitStore prefixes
+_STORE_PREFIX = re.compile('(?<![a-zA-Z0-9_])store[0-9]+_')
+
 
 class UnitError(Exception):
     """Base class for errors relating to calculating units."""
@@ -186,7 +189,7 @@ class UnitStore(object):
         # same registry.
         self._id = UnitStore._next_id
         UnitStore._next_id += 1
-        self._prefix = 's' + str(self._id) + '_'
+        self._prefix = 'store' + str(self._id) + '_'
 
         # Get unit registry
         if store is None:
@@ -280,6 +283,22 @@ class UnitStore(object):
         assert isinstance(unit, self._registry.Unit)
         return quantity.to(unit)
 
+    def format(self, unit, base_units=False):
+        """
+        Returns a string representation of a unit.
+
+        :param unit: The ``Unit`` object to format.
+        :param base_units: If this optional argument is set to ``True`` the string will show the base units expansion of
+            the given unit.
+        """
+        if base_units:
+            base = self._registry.get_base_units(unit)
+            text = str(base[0]) + ' ' + str(base[1])
+            # Remove unit store prefixes from any unit store
+            return _STORE_PREFIX.sub('', text)
+        else:
+            return _STORE_PREFIX.sub('', str(unit))
+
     def get_unit(self, name):
         """Retrieves the ``Unit`` object with the given name.
 
@@ -349,14 +368,6 @@ class UnitStore(object):
 
         return self._prefix + name
 
-    def show_base_units(self, unit):
-        """
-        Returns a string show the base units of the given unit.
-        """
-        base = self._registry.get_base_units(unit)
-        text = str(base[0]) + ' ' + str(base[1])
-        return text.replace(self._prefix, '')
-
 
 class UnitCalculator(object):
     """


Test patch:
diff --git a/tests/test_parser.py b/tests/test_parser.py
index b6ade76..f7232ad 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -218,7 +218,7 @@ class TestParser(object):
         p = parser.Parser(cellml)
         m = p.parse()
         i = m.units.get_unit('inch')
-        assert m.units.show_base_units(i) == '0.0254 meter'
+        assert m.units.format(i, True) == '0.0254 meter'
 
     def test_undefined_variable(self):
         """ Tests parser exception for undefined variable. """
diff --git a/tests/test_units.py b/tests/test_units.py
index e2c179d..551406a 100644
--- a/tests/test_units.py
+++ b/tests/test_units.py
@@ -57,9 +57,9 @@ class TestUnits(object):
 
         unitstore = UnitStore()
         assert not unitstore.is_defined('uu')
-        unitstore.add_base_unit('uu')
+        uu = unitstore.add_base_unit('uu')
         assert unitstore.is_defined('uu')
-        assert unitstore.show_base_units(unitstore.get_unit('uu')) == '1.0 uu'
+        assert unitstore.format(uu, True) == '1.0 uu'
 
         # Duplicate unit definition
         with pytest.raises(ValueError):
@@ -78,6 +78,25 @@ class TestUnits(object):
         y = store.convert(x, mm)
         assert y == 5000 * mm
 
+    def test_format(self):
+        """Tests formatting of units."""
+
+        # Test basic formatting and base unit expansion
+        store = UnitStore()
+        store.add_unit('x', 'second')
+        y = store.add_base_unit('y')
+        z = store.add_unit('z', 'y * meter / x ** 2')
+        assert store.format(z) == 'z'
+        assert store.format(z, True) == '1.0 meter * y / second ** 2'
+
+        # Test if internal name manipulation can mess things up
+        a = store.add_base_unit(str(y))
+        b = store.add_unit(str(z), '2 * z /' + str(y))
+        assert store.format(a) == str(y)
+        assert store.format(b) == str(z)
+        assert store.format(a, True) == '1.0 ' + str(y)
+        assert store.format(b, True) == '2.0 meter * y / second ** 2 / ' + str(y)
+
     def test_get_unit(self):
         """Tests UnitStore.get_unit()."""