Merge 9244ea2 into 16cc9c8

cf_units/_udunits2_parser/graph.py

@@ -109,3 +109,32 @@ class Timestamp(Terminal):
     # This is likely to change in the future, but there are some
     # gnarly test cases, and should not be undertaken lightly.
     pass
+
+
+class Visitor:
+    """
+    This class may be used to help traversing an expression graph.
+
+    It follows the same pattern as the Python ``ast.NodeVisitor``.
+    Users should typically not need to override either ``visit`` or
+    ``generic_visit``, and should instead implement ``visit_<ClassName>``.
+
+    This class is used in cf_units.latex to generate a latex representation
+    of an expression graph.
+
+    """
+    def visit(self, node):
+        """Visit a node."""
+        method = 'visit_' + node.__class__.__name__
+        visitor = getattr(self, method, self.generic_visit)
+        return visitor(node)
+
+    def generic_visit(self, node):
+        """
+        Called if no explicit visitor function exists for a node.
+
+        Can also be called by ``visit_<ClassName>`` implementations
+        if children of the node are to be processed.
+
+        """
+        return [self.visit(child) for child in node.children()]

cf_units/conftest.py

@@ -40,9 +40,14 @@
     parse_test_files = glob.glob(
         os.path.join(here, 'tests', 'integration', 'parse', '*.py'))
 
+    # Files that are python3 only.
+    python3_specific = [os.path.join(here, 'tex.py'),
+                        os.path.join(here, 'tests', 'test_tex.py')]
+
     # collect_ignore is the special variable that pytest reads to
     # indicate which files should be ignored (and not even imported).
     # See also https://docs.pytest.org/en/latest/example/pythoncollection.html
     collect_ignore = (list(all_parse_py) +
                       list(all_compiled_parse_py) +
-                      list(parse_test_files))
+                      list(parse_test_files) +
+                      python3_specific)

cf_units/tests/test_tex.py

@@ -0,0 +1,61 @@
+# (C) British Crown Copyright 2019, Met Office
+#
+# This file is part of cf-units.
+#
+# cf-units is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# cf-units is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with cf-units.  If not, see <http://www.gnu.org/licenses/>.
+
+import pytest
+import six
+
+from cf_units.tex import tex
+
+
+if six.PY2:
+    pytest.skip('skipping latex in Python2')
+
+
+def test_basic():
+    u = 'kg kg-1'
+    assert tex(u) == r'{kg}\cdot{{kg}^{-1}}'
+
+
+def test_identifier_micro():
+    u = 'microW m-2'
+    assert tex(u) == r'{{\mu}W}\cdot{{m}^{-2}}'
+
+
+def test_raise():
+    u = 'm^2'
+    assert tex(u) == r'{m}^{2}'
+
+
+def test_multiply():
+    u = 'foo bar'
+    assert tex(u) == r'{foo}\cdot{bar}'
+
+
+def test_divide():
+    u = 'foo per bar'
+    assert tex(u) == r'\frac{foo}{bar}'
+
+
+def test_shift():
+    u = 'foo @ 50'
+    assert tex(u) == r'{foo} @ {50}'
+
+
+def test_complex():
+    u = 'microW^2 / (5^-2)π per sec @ 42'
+    expected = r'{\frac{{\frac{{{\mu}W}^{2}}{{5}^{-2}}}\cdot{π}}{sec}} @ {42}'
+    assert tex(u) == expected

cf_units/tex.py

@@ -0,0 +1,58 @@
+# (C) British Crown Copyright 2019, Met Office
+#
+# This file is part of cf-units.
+#
+# cf-units is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# cf-units is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with cf-units.  If not, see <http://www.gnu.org/licenses/>.
+
+import six
+
+if six.PY2:
+    raise ImportError('Python3 required for cf-units latex support')
+
+import cf_units._udunits2_parser.graph as graph  # noqa: E402
+from cf_units._udunits2_parser import parse as _parse  # noqa: E402
+
+
+class TeXVisitor(graph.Visitor):
+    def _format(self, fmt, lhs, rhs):
+        return fmt.format(self.visit(lhs), self.visit(rhs))
+
+    def visit_Identifier(self, node):
+        token = str(node)
+        if token.startswith('micro'):
+            token = token.replace('micro', '{\mu}')
+        return token
+
+    def visit_Raise(self, node):
+        return self._format('{{{}}}^{{{}}}', node.lhs, node.rhs)
+
+    def visit_Multiply(self, node):
+        return self._format(r'{{{}}}\cdot{{{}}}', node.lhs, node.rhs)
+
+    def visit_Divide(self, node):
+        return self._format(r'\frac{{{}}}{{{}}}', node.lhs, node.rhs)
+
+    def visit_Shift(self, node):
+        return self._format('{{{}}} @ {{{}}}', node.unit, node.shift_from)
+
+    def generic_visit(self, node):
+        result = [self.visit(child) for child in node.children()]
+        if not result:
+            result = str(node)
+        return result
+
+
+def tex(unit_str):
+    tree = _parse(unit_str)
+    return TeXVisitor().visit(tree)