Add q2str to convert quaternion to string (bdaiinstitute#158)

petercorke · web-flow · commit a08f25ed97d4 · 2025-01-26T01:02:17.000-05:00
diff --git a/spatialmath/base/__init__.py b/spatialmath/base/__init__.py
@@ -208,6 +208,7 @@
     "qdotb",
     "qangle",
     "qprint",
+    "q2str",
     # spatialmath.base.transforms2d
     "rot2",
     "trot2",
diff --git a/spatialmath/base/quaternions.py b/spatialmath/base/quaternions.py
@@ -19,9 +19,11 @@
 import scipy.interpolate as interpolate
 from typing import Optional
 from functools import lru_cache
+import warnings
 
 _eps = np.finfo(np.float64).eps
 
+
 def qeye() -> QuaternionArray:
     """
     Create an identity quaternion
@@ -56,7 +58,7 @@ def qpure(v: ArrayLike3) -> QuaternionArray:
 
     .. runblock:: pycon
 
-        >>> from spatialmath.base import pure, qprint
+        >>> from spatialmath.base import qpure, qprint
         >>> q = qpure([1, 2, 3])
         >>> qprint(q)
     """
@@ -1088,14 +1090,53 @@ def qangle(q1: ArrayLike4, q2: ArrayLike4) -> float:
     return 4.0 * math.atan2(smb.norm(q1 - q2), smb.norm(q1 + q2))
 
 
+def q2str(
+    q: Union[ArrayLike4, ArrayLike4],
+    delim: Optional[Tuple[str, str]] = ("<", ">"),
+    fmt: Optional[str] = "{: .4f}",
+) -> str:
+    """
+    Format a quaternion as a string
+
+    :arg q: unit-quaternion
+    :type q: array_like(4)
+    :arg delim: 2-list of delimeters [default ('<', '>')]
+    :type delim: list or tuple of strings
+    :arg fmt: printf-style format soecifier [default '{: .4f}']
+    :type fmt: str
+    :return: formatted string
+    :rtype: str
+
+    Format the quaternion in a human-readable form as::
+
+        S  D1  VX VY VZ D2
+
+    where S, VX, VY, VZ are the quaternion elements, and D1 and D2 are a pair
+    of delimeters given by `delim`.
+
+    .. runblock:: pycon
+
+        >>> from spatialmath.base import q2str, qrand
+        >>> q = [1, 2, 3, 4]
+        >>> q2str(q)
+        >>> q = qrand()   # a unit quaternion
+        >>> q2str(q, delim=('<<', '>>'))
+
+    :seealso: :meth:`qprint`
+    """
+    q = smb.getvector(q, 4)
+    template = "# {} #, #, # {}".replace("#", fmt)
+    return template.format(q[0], delim[0], q[1], q[2], q[3], delim[1])
+
+
 def qprint(
     q: Union[ArrayLike4, ArrayLike4],
     delim: Optional[Tuple[str, str]] = ("<", ">"),
     fmt: Optional[str] = "{: .4f}",
     file: Optional[TextIO] = sys.stdout,
-) -> str:
+) -> None:
     """
-    Format a quaternion
+    Format a quaternion to a file
 
     :arg q: unit-quaternion
     :type q: array_like(4)
@@ -1105,8 +1146,6 @@ def qprint(
     :type fmt: str
     :arg file: destination for formatted string [default sys.stdout]
     :type file: file object
-    :return: formatted string
-    :rtype: str
 
     Format the quaternion in a human-readable form as::
 
@@ -1117,23 +1156,23 @@ def qprint(
 
     By default the string is written to `sys.stdout`.
 
-    If `file=None` then a string is returned.
-
     .. runblock:: pycon
 
         >>> from spatialmath.base import qprint, qrand
         >>> q = [1, 2, 3, 4]
         >>> qprint(q)
         >>> q = qrand()   # a unit quaternion
         >>> qprint(q, delim=('<<', '>>'))
+
+    :seealso: :meth:`q2str`
     """
     q = smb.getvector(q, 4)
-    template = "# {} #, #, # {}".replace("#", fmt)
-    s = template.format(q[0], delim[0], q[1], q[2], q[3], delim[1])
-    if file:
-        file.write(s + "\n")
-    else:
-        return s
+    if file is None:
+        warnings.warn(
+            "Usage: qprint(..., file=None) -> str is deprecated, use q2str() instead",
+            DeprecationWarning,
+        )
+    print(q2str(q, delim=delim, fmt=fmt), file=file)
 
 
 if __name__ == "__main__":  # pragma: no cover
diff --git a/spatialmath/quaternion.py b/spatialmath/quaternion.py
@@ -920,7 +920,7 @@ def __str__(self) -> str:
             delim = ("<<", ">>")
         else:
             delim = ("<", ">")
-        return "\n".join([smb.qprint(q, file=None, delim=delim) for q in self.data])
+        return "\n".join([smb.q2str(q, delim=delim) for q in self.data])
 
 
 # ========================================================================= #
diff --git a/tests/base/test_quaternions.py b/tests/base/test_quaternions.py
@@ -36,6 +36,7 @@
 import spatialmath.base as tr
 from spatialmath.base.quaternions import *
 import spatialmath as sm
+import io
 
 
 class TestQuaternion(unittest.TestCase):
@@ -96,19 +97,32 @@ def test_ops(self):
             ),
             True,
         )
+        nt.assert_equal(isunitvec(qrand()), True)
 
-        s = qprint(np.r_[1, 1, 0, 0], file=None)
-        nt.assert_equal(isinstance(s, str), True)
-        nt.assert_equal(len(s) > 2, True)
-        s = qprint([1, 1, 0, 0], file=None)
+    def test_display(self):
+        s = q2str(np.r_[1, 2, 3, 4])
         nt.assert_equal(isinstance(s, str), True)
-        nt.assert_equal(len(s) > 2, True)
+        nt.assert_equal(s, " 1.0000 <  2.0000,  3.0000,  4.0000 >")
+
+        s = q2str([1, 2, 3, 4])
+        nt.assert_equal(s, " 1.0000 <  2.0000,  3.0000,  4.0000 >")
 
+        s = q2str([1, 2, 3, 4], delim=("<<", ">>"))
+        nt.assert_equal(s, " 1.0000 <<  2.0000,  3.0000,  4.0000 >>")
+
+        s = q2str([1, 2, 3, 4], fmt="{:20.6f}")
         nt.assert_equal(
-            qprint([1, 2, 3, 4], file=None), " 1.0000 <  2.0000,  3.0000,  4.0000 >"
+            s,
+            "            1.000000 <             2.000000,             3.000000,             4.000000 >",
         )
 
-        nt.assert_equal(isunitvec(qrand()), True)
+        # would be nicer to do this with redirect_stdout() from contextlib but that
+        # fails because file=sys.stdout is maybe assigned at compile time, so when
+        # contextlib changes sys.stdout, qprint() doesn't see it
+
+        f = io.StringIO()
+        qprint(np.r_[1, 2, 3, 4], file=f)
+        nt.assert_equal(f.getvalue().rstrip(), " 1.0000 <  2.0000,  3.0000,  4.0000 >")
 
     def test_rotation(self):
         # rotation matrix to quaternion
@@ -227,12 +241,12 @@ def test_r2q(self):
 
     def test_qangle(self):
         # Test function that calculates angle between quaternions
-        q1 = [1., 0, 0, 0]
-        q2 = [1 / np.sqrt(2), 0, 1 / np.sqrt(2), 0]   # 90deg rotation about y-axis
+        q1 = [1.0, 0, 0, 0]
+        q2 = [1 / np.sqrt(2), 0, 1 / np.sqrt(2), 0]  # 90deg rotation about y-axis
         nt.assert_almost_equal(qangle(q1, q2), np.pi / 2)
 
-        q1 = [1., 0, 0, 0]
-        q2 = [1 / np.sqrt(2), 1 / np.sqrt(2), 0, 0]   # 90deg rotation about x-axis
+        q1 = [1.0, 0, 0, 0]
+        q2 = [1 / np.sqrt(2), 1 / np.sqrt(2), 0, 0]  # 90deg rotation about x-axis
         nt.assert_almost_equal(qangle(q1, q2), np.pi / 2)
 
 
