Support substitution

.gitignore

@@ -90,7 +90,7 @@ ipython_config.py
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
-# .python-version
+.python-version
 
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.

pyproject.toml

@@ -42,6 +42,7 @@ dev = [
     "build>=0.8",
     "black>=22.10",
     "flake8>=5.0",
+    "notebook>=6.5.1",
     "pyproject-flake8>=5.0",
     "pytest>=7.1",
     "twine>=4.0",

src/latexify/latexify_visitor.py

@@ -305,6 +305,43 @@ def visit_GeneratorExp_set_bounds(self, node):  # pylint: disable=invalid-name
             "but {} were given".format(len(comprehensions))
         )
 
+    # Until 3.8
+    def visit_Index(self, node: ast.Index, action) -> str:
+        """Visitor for the Index nodes."""
+        return self.visit(node.value)
+
+    def convert_nested_subscripts(self, node: ast.Subscript) -> tuple[str, list[str]]:
+        """Helper function to convert nested subscription.
+
+        This function converts x[i][j][...] to "x" and ["i", "j", ...]
+
+        Args:
+            node: ast.Subscript node to be converted.
+
+        Returns:
+            Tuple of following strings:
+                - The root value of the subscription.
+                - Sequence of incices.
+        """
+        if isinstance(node.value, ast.Subscript):
+            value, indices = self.convert_nested_subscripts(node.value)
+        else:
+            value = self.visit(node.value)
+            indices = []
+
+        indices.append(self.visit(node.slice))
+        return value, indices
+
+    def visit_Subscript(self, node: ast.Subscript, action) -> str:
+        """Visitor of the Subscript nodes."""
+        value, indices = self.convert_nested_subscripts(node)
+
+        # TODO(odashi):
+        # "[i][j][...]" may be a possible representation as well as "i, j. ..."
+        indices_str = "{" + ", ".join(indices) + "}"
+
+        return f"{{{value}_{indices_str}}}"
+
     def visit_comprehension_set_bounds(self, node):  # pylint: disable=invalid-name
         """Visit a comprehension node, which represents a for clause"""
         var = self.visit(node.target)

src/latexify/latexify_visitor_test.py

@@ -65,3 +65,19 @@ def test_visit_boolop(code: str, latex: str) -> None:
     tree = ast.parse(code).body[0].value
     assert isinstance(tree, ast.BoolOp)
     assert LatexifyVisitor().visit(tree) == latex
+
+
+@pytest.mark.parametrize(
+    "code,latex",
+    [
+        ("x[0]", "{x_{0}}"),
+        ("x[0][1]", "{x_{0, 1}}"),
+        ("x[0][1][2]", "{x_{0, 1, 2}}"),
+        ("x[foo]", "{x_{foo}}"),
+        ("x[math.floor(x)]", r"{x_{\left\lfloor{x}\right\rfloor}}"),
+    ],
+)
+def test_visit_subscript(code: str, latex: str) -> None:
+    tree = ast.parse(code).body[0].value
+    assert isinstance(tree, ast.Subscript)
+    assert LatexifyVisitor().visit(tree) == latex