add number formatting options

Author: dbatten5

src/arraytex/api.py

@@ -1,28 +1,51 @@
 """Main package API."""
-
+import re
 from typing import Any
+from typing import Optional
+from typing import Union
 
 import numpy as np
 from numpy.typing import NDArray
 
 
-def to_matrix(arr: NDArray[Any]) -> str:
+def to_matrix(
+    arr: NDArray[Any],
+    style: str = "b",
+    num_format: Optional[str] = None,
+) -> str:
     """Convert a numpy.NDArray to LaTeX matrix.
 
     Args:
         arr: the array to be converted
+        style: a style formatter string, either "b" for "bmatrid" or "p" for "pmatrix"
+        num_format: a number formatter string, e.g. ".2f"
 
     Returns:
         the string representation of the array
     """
-    environment = "bmatrix"
+
+    def num_formatter(x: Union[int, float]) -> str:
+        return f"%{num_format}" % x
+
+    environment = f"{style}matrix"
+
+    formatter = {}
+    if num_format:
+        formatter.update({"float_kind": num_formatter, "int_kind": num_formatter})
+
     lines = (
-        np.array2string(arr, max_line_width=np.inf)  # type: ignore  # noqa
+        np.array2string(arr, max_line_width=np.inf, formatter=formatter)  # type: ignore  # noqa
         .replace("[", "")
         .replace("]", "")
-        .splitlines()
     )
+
+    if num_format and "e" in num_format:
+        pattern = r"e(-?\d+)"
+        replace = r"\\mathrm{e}{\g<1>}"
+        lines = re.sub(pattern, replace, lines)
+
     rv = [f"\\begin{{{environment}}}"]
-    rv += ["  " + " & ".join(line.split()) + r" \\" for line in lines]
+    rv += ["  " + " & ".join(line.split()) + r" \\" for line in lines.splitlines()]
     rv += [f"\\end{{{environment}}}"]
+
     return "\n".join(rv)

tests/test_api.py

@@ -7,7 +7,7 @@
 class TestToMatrix:
     """Tests for the `to_matrix` function."""
 
-    def test_success(self) -> None:
+    def test_default(self) -> None:
         """A formatted matrix is returned."""
         mat = np.array(
             [
@@ -23,5 +23,62 @@ def test_success(self) -> None:
             == r"""\begin{bmatrix}
   1 & 2 & 3 \\
   4 & 5 & 6 \\
+\end{bmatrix}"""
+        )
+
+    def test_style(self) -> None:
+        """An arg can be supplied to determine the style of the returned matrix."""
+        mat = np.array(
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+            ]
+        )
+
+        out = to_matrix(mat, style="p")
+
+        assert (
+            out
+            == r"""\begin{pmatrix}
+  1 & 2 & 3 \\
+  4 & 5 & 6 \\
+\end{pmatrix}"""
+        )
+
+    def test_decimal_format(self) -> None:
+        """Decimal style of number formatting can be selected."""
+        mat = np.array(
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+            ]
+        )
+
+        out = to_matrix(mat, num_format=".2f")
+
+        assert (
+            out
+            == r"""\begin{bmatrix}
+  1.00 & 2.00 & 3.00 \\
+  4.00 & 5.00 & 6.00 \\
+\end{bmatrix}"""
+        )
+
+    def test_scientific_notation(self) -> None:
+        """Scientific notation can be selected."""
+        mat = np.array(
+            [
+                [0.001, 0.002, 0.003],
+                [0.004, 0.005, 0.006],
+            ]
+        )
+
+        out = to_matrix(mat, num_format=".2e")
+
+        assert (
+            out
+            == r"""\begin{bmatrix}
+  1.00\mathrm{e}{-03} & 2.00\mathrm{e}{-03} & 3.00\mathrm{e}{-03} \\
+  4.00\mathrm{e}{-03} & 5.00\mathrm{e}{-03} & 6.00\mathrm{e}{-03} \\
 \end{bmatrix}"""
         )