Repository: google/latexify_py

Issue:
NDArray support
Sometimes users want to support NDArrays in this library, e.g., `np.array([[a, b], [c, d]])` to 

$$\left( \begin{array}{cc}a&b \\\\ c&d\end{array} \right)$$

This is definitely useful, but often it is not possible to compile:

- Large arrays.
- Arrays with more than 2 dimensions.

I think we could start implementing it with reasonable restrictions (e.g., only NDArrays with <= 2 dimensions with small lengths), but it would be also good to keep discussion.

Refs:
- #38
- #53
- #75

Comments:

Patch:
diff --git a/src/latexify/codegen/function_codegen.py b/src/latexify/codegen/function_codegen.py
index 9ef24a0..cf83122 100644
--- a/src/latexify/codegen/function_codegen.py
+++ b/src/latexify/codegen/function_codegen.py
@@ -334,18 +334,84 @@ class FunctionCodegen(ast.NodeVisitor):
         wrapped = [r"\mathopen{}\left( " + s + r" \mathclose{}\right)" for s in conds]
         return r" \land ".join(wrapped)
 
+    def _generate_sum_prod(self, node: ast.Call) -> str | None:
+        """Generates sum/prod expression.
+
+        Args:
+            node: ast.Call node containing the sum/prod invocation.
+
+        Returns:
+            Generated LaTeX, or None if the node has unsupported syntax.
+        """
+        if not isinstance(node.args[0], ast.GeneratorExp):
+            return None
+
+        name = ast_utils.extract_function_name_or_none(node)
+        assert name is not None
+
+        elt, scripts = self._get_sum_prod_info(node.args[0])
+        scripts_str = [rf"\{name}_{{{lo}}}^{{{up}}}" for lo, up in scripts]
+        return (
+            " ".join(scripts_str)
+            + rf" \mathopen{{}}\left({{{elt}}}\mathclose{{}}\right)"
+        )
+
+    def _generate_matrix(self, node: ast.Call) -> str | None:
+        """Generates matrix expression.
+
+        Args:
+            node: ast.Call node containing the ndarray invocation.
+
+        Returns:
+            Generated LaTeX, or None if the node has unsupported syntax.
+        """
+
+        def generate_matrix_from_array(data: list[list[str]]) -> str:
+            """Helper to generate a bmatrix environment."""
+            contents = r" \\ ".join(" & ".join(row) for row in data)
+            return r"\begin{bmatrix} " + contents + r" \end{bmatrix}"
+
+        arg = node.args[0]
+        if not isinstance(arg, ast.List) or not arg.elts:
+            # Not an array or no rows
+            return None
+
+        row0 = arg.elts[0]
+
+        if not isinstance(row0, ast.List):
+            # Maybe 1 x N array
+            return generate_matrix_from_array([[self.visit(x) for x in arg.elts]])
+
+        if not row0.elts:
+            # No columns
+            return None
+
+        ncols = len(row0.elts)
+
+        if not all(
+            isinstance(row, ast.List) and len(row.elts) == ncols for row in arg.elts
+        ):
+            # Length mismatch
+            return None
+
+        return generate_matrix_from_array(
+            [[self.visit(x) for x in row.elts] for row in arg.elts]
+        )
+
     def visit_Call(self, node: ast.Call) -> str:
         """Visit a call node."""
         func_name = ast_utils.extract_function_name_or_none(node)
 
-        # Special processing for sum and prod.
-        if func_name in ("sum", "prod") and isinstance(node.args[0], ast.GeneratorExp):
-            elt, scripts = self._get_sum_prod_info(node.args[0])
-            scripts_str = [rf"\{func_name}_{{{lo}}}^{{{up}}}" for lo, up in scripts]
-            return (
-                " ".join(scripts_str)
-                + rf" \mathopen{{}}\left({{{elt}}}\mathclose{{}}\right)"
-            )
+        # Special treatments for some functions.
+        if func_name in ("sum", "prod"):
+            special_latex = self._generate_sum_prod(node)
+        elif func_name in ("array", "ndarray"):
+            special_latex = self._generate_matrix(node)
+        else:
+            special_latex = None
+
+        if special_latex is not None:
+            return special_latex
 
         # Function signature (possibly an expression).
         default_func_str = self.visit(node.func)


Test patch:
diff --git a/src/latexify/codegen/function_codegen_test.py b/src/latexify/codegen/function_codegen_test.py
index 2dc880a..d7628be 100644
--- a/src/latexify/codegen/function_codegen_test.py
+++ b/src/latexify/codegen/function_codegen_test.py
@@ -744,3 +744,46 @@ def test_use_set_symbols_compare(code: str, latex: str) -> None:
     tree = ast.parse(code).body[0].value
     assert isinstance(tree, ast.Compare)
     assert function_codegen.FunctionCodegen(use_set_symbols=True).visit(tree) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        ("array(1)", r"\mathrm{array}\mathopen{}\left({1}\mathclose{}\right)"),
+        (
+            "array([])",
+            r"\mathrm{array}\mathopen{}\left(\left[ \right] \mathclose{}\right)",
+        ),
+        ("array([1])", r"\begin{bmatrix} {1} \end{bmatrix}"),
+        ("array([1, 2, 3])", r"\begin{bmatrix} {1} & {2} & {3} \end{bmatrix}"),
+        (
+            "array([[]])",
+            r"\mathrm{array}\mathopen{}\left("
+            r"\left[ \left[ \right] \right] "
+            r"\mathclose{}\right)",
+        ),
+        ("array([[1]])", r"\begin{bmatrix} {1} \end{bmatrix}"),
+        ("array([[1], [2], [3]])", r"\begin{bmatrix} {1} \\ {2} \\ {3} \end{bmatrix}"),
+        (
+            "array([[1], [2], [3, 4]])",
+            r"\mathrm{array}\mathopen{}\left("
+            r"\left[ "
+            r"\left[ {1}\right] \space,\space "
+            r"\left[ {2}\right] \space,\space "
+            r"\left[ {3}\space,\space {4}\right] "
+            r"\right] "
+            r"\mathclose{}\right)",
+        ),
+        (
+            "array([[1, 2], [3, 4], [5, 6]])",
+            r"\begin{bmatrix} {1} & {2} \\ {3} & {4} \\ {5} & {6} \end{bmatrix}",
+        ),
+        # Only checks two cases for ndarray.
+        ("ndarray(1)", r"\mathrm{ndarray}\mathopen{}\left({1}\mathclose{}\right)"),
+        ("ndarray([1])", r"\begin{bmatrix} {1} \end{bmatrix}"),
+    ],
+)
+def test_numpy_array(code: str, latex: str) -> None:
+    tree = ast.parse(code).body[0].value
+    assert isinstance(tree, ast.Call)
+    assert function_codegen.FunctionCodegen().visit(tree) == latex