add to_matrix method

Author: dbatten5

src/arraytex/api.py

@@ -0,0 +1,28 @@
+"""Main package API."""
+
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def to_matrix(arr: NDArray[Any]) -> str:
+    """Convert a numpy.NDArray to LaTeX matrix.
+
+    Args:
+        arr: the array to be converted
+
+    Returns:
+        the string representation of the array
+    """
+    environment = "bmatrix"
+    lines = (
+        np.array2string(arr, max_line_width=np.inf)  # type: ignore  # noqa
+        .replace("[", "")
+        .replace("]", "")
+        .splitlines()
+    )
+    rv = [f"\\begin{{{environment}}}"]
+    rv += ["  " + " & ".join(line.split()) + r" \\" for line in lines]
+    rv += [f"\\end{{{environment}}}"]
+    return "\n".join(rv)

tests/test_api.py

@@ -0,0 +1,27 @@
+"""Tests for the main API."""
+import numpy as np
+
+from arraytex.api import to_matrix
+
+
+class TestToMatrix:
+    """Tests for the `to_matrix` function."""
+
+    def test_success(self) -> None:
+        """A formatted matrix is returned."""
+        mat = np.array(
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+            ]
+        )
+
+        out = to_matrix(mat)
+
+        assert (
+            out
+            == r"""\begin{bmatrix}
+  1 & 2 & 3 \\
+  4 & 5 & 6 \\
+\end{bmatrix}"""
+        )