Enable mypy (#153)

* add mypy

* revert BuiltinFnName

* fix

* fix

* fix

* fix

Author: Yusuke Oda

.github/workflows/ci.yml

@@ -69,3 +69,17 @@ jobs:
           python -m pip install isort
       - name: Check
         run: python -m isort --check src
+  mypy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.x"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install '.[mypy]'
+      - name: Check
+        run: python -m mypy src

checks.sh

@@ -5,3 +5,4 @@ python -m pytest src -vv
 python -m black --check src
 python -m pflake8 src
 python -m isort --check src
+python -m mypy src

pyproject.toml

@@ -45,11 +45,16 @@ dev = [
     "build>=0.8",
     "black>=22.10",
     "flake8>=5.0",
+    "isort>=5.10",
+    "mypy>=0.991",
     "notebook>=6.5.1",
     "pyproject-flake8>=5.0",
     "pytest>=7.1",
     "twine>=4.0",
-    "isort>=5.10",
+]
+mypy = [
+    "mypy>=0.991",
+    "pytest>=7.1",
 ]
 
 [project.urls]

src/latexify/__init__.py

@@ -13,6 +13,3 @@
 
 function = frontend.function
 expression = frontend.expression
-
-# Deprecated
-with_latex = frontend.with_latex

src/latexify/ast_utils_test.py

@@ -121,15 +121,6 @@ def test_extract_int_or_none() -> None:
 
 
 def test_extract_int_or_none_invalid() -> None:
-    # Not a subtree.
-    assert ast_utils.extract_int_or_none(123) is None
-
-    # Not a direct Constant node.
-    assert (
-        ast_utils.extract_int_or_none(ast.Expr(value=ast_utils.make_constant(123)))
-        is None
-    )
-
     # Not a Constant node with int.
     assert ast_utils.extract_int_or_none(ast_utils.make_constant(None)) is None
     assert ast_utils.extract_int_or_none(ast_utils.make_constant(True)) is None
@@ -147,14 +138,6 @@ def test_extract_int() -> None:
 
 
 def test_extract_int_invalid() -> None:
-    # Not a subtree.
-    with pytest.raises(ValueError, match=r"^Unsupported node to extract int"):
-        ast_utils.extract_int(123)
-
-    # Not a direct Constant node.
-    with pytest.raises(ValueError, match=r"^Unsupported node to extract int"):
-        ast_utils.extract_int(ast.Expr(value=ast_utils.make_constant(123)))
-
     # Not a Constant node with int.
     with pytest.raises(ValueError, match=r"^Unsupported node to extract int"):
         ast_utils.extract_int(ast_utils.make_constant(None))

src/latexify/codegen/function_codegen.py

@@ -296,7 +296,11 @@ def visit_Assign(self, node: ast.Assign) -> str:
         return " = ".join(operands)
 
     def visit_Return(self, node: ast.Return) -> str:
-        return self.visit(node.value)
+        return (
+            self.visit(node.value)
+            if node.value is not None
+            else self._convert_constant(None)
+        )
 
     def visit_Tuple(self, node: ast.Tuple) -> str:
         elts = [self.visit(i) for i in node.elts]
@@ -401,15 +405,16 @@ def generate_matrix_from_array(data: list[list[str]]) -> str:
 
         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
+        rows: list[list[str]] = []
 
-        return generate_matrix_from_array(
-            [[self.visit(x) for x in row.elts] for row in arg.elts]
-        )
+        for row in arg.elts:
+            if not isinstance(row, ast.List) or len(row.elts) != ncols:
+                # Length mismatch
+                return None
+
+            rows.append([self.visit(x) for x in row.elts])
+
+        return generate_matrix_from_array(rows)
 
     def visit_Call(self, node: ast.Call) -> str:
         """Visit a call node."""
@@ -427,7 +432,8 @@ def visit_Call(self, node: ast.Call) -> str:
             return special_latex
 
         # Obtains the codegen rule.
-        rule = constants.BUILTIN_FUNCS.get(func_name)
+        rule = constants.BUILTIN_FUNCS.get(func_name) if func_name is not None else None
+
         if rule is None:
             rule = constants.FunctionRule(self.visit(node.func))
 
@@ -556,8 +562,11 @@ def _wrap_binop_operand(
             return self.visit(child)
 
         if isinstance(child, ast.Call):
-            rule = constants.BUILTIN_FUNCS.get(
-                ast_utils.extract_function_name_or_none(child)
+            child_fn_name = ast_utils.extract_function_name_or_none(child)
+            rule = (
+                constants.BUILTIN_FUNCS.get(child_fn_name)
+                if child_fn_name is not None
+                else None
             )
             if rule is not None and rule.is_wrapped:
                 return self.visit(child)
@@ -612,30 +621,35 @@ def visit_If(self, node: ast.If) -> str:
         """Visit an if node."""
         latex = r"\left\{ \begin{array}{ll} "
 
-        while isinstance(node, ast.If):
-            if len(node.body) != 1 or len(node.orelse) != 1:
+        current_stmt: ast.stmt = node
+
+        while isinstance(current_stmt, ast.If):
+            if len(current_stmt.body) != 1 or len(current_stmt.orelse) != 1:
                 raise exceptions.LatexifySyntaxError(
                     "Multiple statements are not supported in If nodes."
                 )
 
-            cond_latex = self.visit(node.test)
-            true_latex = self.visit(node.body[0])
+            cond_latex = self.visit(current_stmt.test)
+            true_latex = self.visit(current_stmt.body[0])
             latex += true_latex + r", & \mathrm{if} \ " + cond_latex + r" \\ "
-            node = node.orelse[0]
+            current_stmt = current_stmt.orelse[0]
 
-        latex += self.visit(node)
+        latex += self.visit(current_stmt)
         return latex + r", & \mathrm{otherwise} \end{array} \right."
 
     def visit_IfExp(self, node: ast.IfExp) -> str:
         """Visit an ifexp node"""
         latex = r"\left\{ \begin{array}{ll} "
-        while isinstance(node, ast.IfExp):
-            cond_latex = self.visit(node.test)
-            true_latex = self.visit(node.body)
+
+        current_expr: ast.expr = node
+
+        while isinstance(current_expr, ast.IfExp):
+            cond_latex = self.visit(current_expr.test)
+            true_latex = self.visit(current_expr.body)
             latex += true_latex + r", & \mathrm{if} \ " + cond_latex + r" \\ "
-            node = node.orelse
+            current_expr = current_expr.orelse
 
-        latex += self.visit(node)
+        latex += self.visit(current_expr)
         return latex + r", & \mathrm{otherwise} \end{array} \right."
 
     def _reduce_stop_parameter(self, node: ast.expr) -> ast.expr:
@@ -768,7 +782,7 @@ def _get_sum_prod_info(
     # Until 3.8
     def visit_Index(self, node: ast.Index) -> str:
         """Visitor for the Index nodes."""
-        return self.visit(node.value)
+        return self.visit(node.value)  # type: ignore[attr-defined]
 
     def _convert_nested_subscripts(self, node: ast.Subscript) -> tuple[str, list[str]]:
         """Helper function to convert nested subscription.

src/latexify/codegen/latex.py

@@ -21,7 +21,7 @@ def __init__(self, raw: str) -> None:
         """
         self._raw = raw
 
-    def __eq__(self, other: object) -> None:
+    def __eq__(self, other: object) -> bool:
         """Checks equality.
 
         Args:

src/latexify/frontend.py

@@ -2,9 +2,8 @@
 
 from __future__ import annotations
 
-import warnings
 from collections.abc import Callable
-from typing import Any
+from typing import Any, overload
 
 from latexify import codegen
 from latexify import config as cfg
@@ -119,45 +118,63 @@ def _repr_latex_(self):
         )
 
 
-def function(*args, **kwargs) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
-    """Translate a function into a corresponding LaTeX representation.
+@overload
+def function(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
+    ...
+
+
+@overload
+def function(**kwargs: Any) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
+    ...
+
+
+def function(
+    fn: Callable[..., Any] | None = None, **kwargs: Any
+) -> LatexifiedFunction | Callable[[Callable[..., Any]], LatexifiedFunction]:
+    """Attach LaTeX pretty-printing to the given function.
 
     This function works with or without specifying the target function as the positional
     argument. The following two syntaxes works similarly.
-        - with_latex(fn, **kwargs)
-        - with_latex(**kwargs)(fn)
+        - latexify.function(fn, **kwargs)
+        - latexify.function(**kwargs)(fn)
 
     Args:
-        *args: No argument, or a callable.
+        fn: Callable to be wrapped.
         **kwargs: Arguments to control behavior. See also get_latex().
 
     Returns:
-        - If the target function is passed directly, returns the wrapped function.
+        - If `fn` is passed, returns the wrapped function.
         - Otherwise, returns the wrapper function with given settings.
     """
-    if len(args) == 1 and isinstance(args[0], Callable):
-        return LatexifiedFunction(args[0], **kwargs)
-
-    def wrapper(fn):
+    if fn is not None:
         return LatexifiedFunction(fn, **kwargs)
 
+    def wrapper(f):
+        return LatexifiedFunction(f, **kwargs)
+
     return wrapper
 
 
-def expression(*args, **kwargs) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
-    """Translate a function into a LaTeX representation without the signature.
+@overload
+def expression(fn: Callable[..., Any], **kwargs: Any) -> LatexifiedFunction:
+    ...
+
+
+@overload
+def expression(**kwargs: Any) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
+    ...
+
+
+def expression(
+    fn: Callable[..., Any] | None = None, **kwargs: Any
+) -> LatexifiedFunction | Callable[[Callable[..., Any]], LatexifiedFunction]:
+    """Attach LaTeX pretty-printing to the given function.
 
     This function is a shortcut for `latexify.function` with the default parameter
     `use_signature=False`.
     """
     kwargs["use_signature"] = kwargs.get("use_signature", False)
-    return function(*args, **kwargs)
-
-
-def with_latex(*args, **kwargs) -> Callable[[Callable[..., Any]], LatexifiedFunction]:
-    """Deprecated. use `latexify.function` instead."""
-    warnings.warn(
-        "`latexify.with_latex` is deprecated. Use `latexify.function` instead.",
-        DeprecationWarning,
-    )
-    return function(*args, **kwargs)
+    if fn is not None:
+        return function(fn, **kwargs)
+    else:
+        return function(**kwargs)

src/latexify/parser.py

@@ -8,12 +8,12 @@
 from collections.abc import Callable
 from typing import Any
 
-import dill
+import dill  # type: ignore[import]
 
 from latexify import exceptions
 
 
-def parse_function(fn: Callable[..., Any]) -> ast.FunctionDef:
+def parse_function(fn: Callable[..., Any]) -> ast.Module:
     """Parses given function.
 
     Args:

src/latexify/transformers/__init__.py

@@ -6,8 +6,8 @@
 from latexify.transformers.prefix_trimmer import PrefixTrimmer
 
 __all__ = [
-    AssignmentReducer,
-    FunctionExpander,
-    IdentifierReplacer,
-    PrefixTrimmer,
+    "AssignmentReducer",
+    "FunctionExpander",
+    "IdentifierReplacer",
+    "PrefixTrimmer",
 ]

src/latexify/transformers/identifier_replacer.py

@@ -5,7 +5,7 @@
 import ast
 import re
 import sys
-from typing import ClassVar
+from typing import ClassVar, cast
 
 
 class IdentifierReplacer(ast.NodeTransformer):
@@ -47,33 +47,32 @@ def _replace_args(self, args: list[ast.arg]) -> list[ast.arg]:
         """Helper function to replace arg names."""
         return [ast.arg(arg=self._mapping.get(a.arg, a.arg)) for a in args]
 
-    def _visit_children(self, children: list[ast.AST]) -> list[ast.AST]:
-        """Helper function to visit all children."""
-        return [self.visit(child) for child in children]
-
     def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
         """Visitor of FunctionDef."""
+
+        visited = cast(ast.FunctionDef, super().generic_visit(node))
+
         if sys.version_info.minor < 8:
             args = ast.arguments(
-                args=self._replace_args(node.args.args),
-                kwonlyargs=self._replace_args(node.args.kwonlyargs),
-                kw_defaults=self._visit_children(node.args.kw_defaults),
-                defaults=self._visit_children(node.args.defaults),
+                args=self._replace_args(visited.args.args),
+                kwonlyargs=self._replace_args(visited.args.kwonlyargs),
+                kw_defaults=visited.args.kw_defaults,
+                defaults=visited.args.defaults,
             )
         else:
             args = ast.arguments(
-                posonlyargs=self._replace_args(node.args.posonlyargs),  # from 3.8
-                args=self._replace_args(node.args.args),
-                kwonlyargs=self._replace_args(node.args.kwonlyargs),
-                kw_defaults=self._visit_children(node.args.kw_defaults),
-                defaults=self._visit_children(node.args.defaults),
+                posonlyargs=self._replace_args(visited.args.posonlyargs),  # from 3.8
+                args=self._replace_args(visited.args.args),
+                kwonlyargs=self._replace_args(visited.args.kwonlyargs),
+                kw_defaults=visited.args.kw_defaults,
+                defaults=visited.args.defaults,
             )
 
         return ast.FunctionDef(
-            name=self._mapping.get(node.name, node.name),
+            name=self._mapping.get(visited.name, visited.name),
             args=args,
-            body=self._visit_children(node.body),
-            decorator_list=self._visit_children(node.decorator_list),
+            body=visited.body,
+            decorator_list=visited.decorator_list,
         )
 
     def visit_Name(self, node: ast.Name) -> ast.Name: