From 6c2d90fca6bf1a5f044f20321436c4075a96489e Mon Sep 17 00:00:00 2001
From: Robert Haschke <rhaschke@users.noreply.github.com>
Date: Tue, 18 Aug 2020 17:41:24 +0200
Subject: [PATCH] Provide support for yaml constructors !degrees and !radians
 (#252)

Co-authored-by: gavanderhoorn <g.a.vanderhoorn@tudelft.nl>
---
 src/xacro/__init__.py       | 21 +++++++++++++++++++++
 test/constructors.yaml      |  3 +++
 test/constructors_bad1.yaml |  1 +
 test/constructors_bad2.yaml |  1 +
 test/test_xacro.py          | 19 +++++++++++++++++++
 5 files changed, 45 insertions(+)
 create mode 100644 test/constructors.yaml
 create mode 100644 test/constructors_bad1.yaml
 create mode 100644 test/constructors_bad2.yaml

diff --git a/src/xacro/__init__.py b/src/xacro/__init__.py
index 47b10de9..2b546ada 100644
--- a/src/xacro/__init__.py
+++ b/src/xacro/__init__.py
@@ -105,9 +105,30 @@ def __getattr__(self, item):
 
     __getitem__ = __getattr__
 
+
+def construct_angle_radians(loader, node):
+    """utility function to construct radian values from yaml"""
+    value = loader.construct_scalar(node).strip()
+    try:
+        return float(eval(value, global_symbols, global_symbols))
+    except SyntaxError as e:
+        raise XacroException("invalid expression: %s" % value)
+
+
+def construct_angle_degrees(loader, node):
+    """utility function for converting degrees into radians from yaml"""
+    value = loader.construct_scalar(node)
+    try:
+        return math.radians(float(value))
+    except ValueError:
+        raise XacroException("invalid degree value: %s" % value)
+
+
 def load_yaml(filename):
     try:
         import yaml
+        yaml.SafeLoader.add_constructor(u'!radians', construct_angle_radians)
+        yaml.SafeLoader.add_constructor(u'!degrees', construct_angle_degrees)
     except Exception:
         raise XacroException("yaml support not available; install python-yaml")
 
diff --git a/test/constructors.yaml b/test/constructors.yaml
new file mode 100644
index 00000000..b018c6aa
--- /dev/null
+++ b/test/constructors.yaml
@@ -0,0 +1,3 @@
+a: !degrees 180
+b: !radians 0.5*pi
+c: 42
diff --git a/test/constructors_bad1.yaml b/test/constructors_bad1.yaml
new file mode 100644
index 00000000..f5309daa
--- /dev/null
+++ b/test/constructors_bad1.yaml
@@ -0,0 +1 @@
+a: !radians os.system('clear')
diff --git a/test/constructors_bad2.yaml b/test/constructors_bad2.yaml
new file mode 100644
index 00000000..3ba9a6ed
--- /dev/null
+++ b/test/constructors_bad2.yaml
@@ -0,0 +1 @@
+a: !radians __import__('os').system('clear')
diff --git a/test/test_xacro.py b/test/test_xacro.py
index 68729ede..0b12c71a 100644
--- a/test/test_xacro.py
+++ b/test/test_xacro.py
@@ -37,6 +37,7 @@
 
 import ast
 from contextlib import contextmanager
+import math
 import os.path
 import re
 import shutil
@@ -1204,6 +1205,24 @@ def test_xacro_exist_optional(self):
         res = '''<a></a>'''
         self.assert_matches(self.quick_xacro(src), res)
 
+    def test_yaml_custom_constructors(self):
+        src = '''
+<a xmlns:xacro="http://www.ros.org/wiki/xacro">
+  <xacro:property name="values" value="${load_yaml('constructors.yaml')}"/>
+  <values a="${values.a}" b="${values.b}" c="${values.c}"/>
+</a>'''
+        res = '''<a><values a="{}" b="{}" c="42"/></a>'''.format(math.pi, 0.5*math.pi)
+        self.assert_matches(self.quick_xacro(src), res)
+
+    def test_yaml_custom_constructors_illegal_expr(self):
+        for file in ['constructors_bad1.yaml', 'constructors_bad2.yaml']:
+          src = '''
+<a xmlns:xacro="http://www.ros.org/wiki/xacro">
+  <xacro:property name="values" value="${{load_yaml({file})}}"/>
+  <values a="${{values.a}}" />
+</a>'''
+        self.assertRaises(xacro.XacroException, self.quick_xacro, src.format(file=file))
+
     def test_macro_default_param_evaluation_order(self):
         src = '''<a xmlns:xacro="http://www.ros.org/wiki/xacro">
 <xacro:macro name="foo" params="arg:=${2*foo}">