Introduce mypy

.github/workflows/python-package.yml

@@ -33,6 +33,9 @@ jobs:
         flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - name: Type check with mypy
+      run: |
+        mypy .
     - name: Test with pytest
       run: |
         pytest

README.md

@@ -49,6 +49,7 @@ python -m atcoder [your-source-code] -o [single-combined-code]
 
 + [flake8](https://pypi.org/project/flake8/)
 + [pep8-naming](https://pypi.org/project/pep8-naming/)
++ [mypy](https://pypi.org/project/mypy/)
 
 ## How to contribute
 

README_ja.md

@@ -46,6 +46,7 @@ python -m atcoder [your-source-code] -o [single-combined-code]
 
 + [flake8](https://pypi.org/project/flake8/)
 + [pep8-naming](https://pypi.org/project/pep8-naming/)
++ [mypy](https://pypi.org/project/mypy/)
 
 ## 本レポジトリに貢献する方法
 

atcoder/__main__.py

@@ -3,7 +3,7 @@
 import importlib
 import inspect
 import re
-from typing import List, Optional
+from typing import List, Optional, cast
 
 
 class ImportInfo:
@@ -34,14 +34,14 @@ def iter_child_nodes(
         for name in node.names:
             if re.match(r'^atcoder\.?', name.name):
                 if hasattr(node, 'end_lineno'):
-                    end_lineno = node.end_lineno
+                    end_lineno = cast(int, node.end_lineno)  # type: ignore
                 else:
                     end_lineno = node.lineno
                 import_info = ImportInfo(node.lineno, end_lineno)
     elif isinstance(node, ast.ImportFrom):
-        if re.match(r'^atcoder\.?', node.module):
+        if re.match(r'^atcoder\.?', cast(str, node.module)):
             if hasattr(node, 'end_lineno'):
-                end_lineno = node.end_lineno
+                end_lineno = cast(int, node.end_lineno)  # type: ignore
             else:
                 end_lineno = node.lineno
             import_info = ImportInfo(node.lineno, end_lineno, node.module)
@@ -53,7 +53,7 @@ def iter_child_nodes(
 
 class ModuleImporter:
     def __init__(self) -> None:
-        self.imported_modules = []
+        self.imported_modules: List[str] = []
 
     def import_module(self, import_from: Optional[str], name: str,
                       asname: Optional[str] = None) -> str:
@@ -79,16 +79,16 @@ def import_module(self, import_from: Optional[str], name: str,
             import_lines = []
             for import_info in imports:
                 result += self.import_module(
-                    import_info.import_from, import_info.name,
+                    import_info.import_from, cast(str, import_info.name),
                     import_info.asname)
                 for line in range(import_info.lineno - 1,
                                   import_info.end_lineno):
                     import_lines.append(line)
 
-            for lineno, line in enumerate(lines):
+            for lineno, line_str in enumerate(lines):
                 if lineno not in import_lines:
                     continue
-                lines[lineno] = '# ' + line  # TODO(not): indent
+                lines[lineno] = '# ' + line_str  # TODO(not): indent
 
             modules = module_name.split('.')
             for i in range(len(modules) - 1):
@@ -104,7 +104,7 @@ def import_module(self, import_from: Optional[str], name: str,
             imported = []
             for import_info in imports:
                 if import_info.import_from is None:
-                    modules = import_info.name.split('.')
+                    modules = cast(str, import_info.name).split('.')
                     for i in range(len(modules)):
                         import_name = '.'.join(modules[:i + 1])
                         if import_name in imported:
@@ -149,14 +149,15 @@ def main() -> None:
     import_lines = []
     for import_info in imports:
         result += importer.import_module(
-            import_info.import_from, import_info.name, import_info.asname)
+            import_info.import_from, cast(str, import_info.name),
+            import_info.asname)
         for line in range(import_info.lineno - 1, import_info.end_lineno):
             import_lines.append(line)
 
-    for lineno, line in enumerate(lines):
+    for lineno, line_str in enumerate(lines):
         if lineno not in import_lines:
             continue
-        lines[lineno] = '# ' + line  # TODO(not): indent
+        lines[lineno] = '# ' + line_str  # TODO(not): indent
     result += ''.join(lines)
 
     if args.output:

atcoder/_scc.py

@@ -30,7 +30,7 @@ class SCCGraph:
 
     def __init__(self, n: int) -> None:
         self._n = n
-        self._edges = []
+        self._edges: typing.List[typing.Tuple[int, int]] = []
 
     def num_vertices(self) -> int:
         return self._n
@@ -94,7 +94,7 @@ def scc(self) -> typing.List[typing.List[int]]:
         counts = [0] * group_num
         for x in ids[1]:
             counts[x] += 1
-        groups = [[] for _ in range(group_num)]
+        groups: typing.List[typing.List[int]] = [[] for _ in range(group_num)]
         for i in range(self._n):
             groups[ids[1][i]].append(i)
 

atcoder/convolution.py

@@ -57,8 +57,8 @@ def _butterfly_inv(a: typing.List[Modint]) -> None:
     h = atcoder._bit._ceil_pow2(n)
 
     if a[0].mod() not in _sum_ie:
-        es = [0] * 30  # es[i]^(2^(2+i)) == 1
-        ies = [0] * 30
+        es = [Modint(0)] * 30  # es[i]^(2^(2+i)) == 1
+        ies = [Modint(0)] * 30
         cnt2 = atcoder._bit._bsf(a[0].mod() - 1)
         e = Modint(g) ** ((a[0].mod() - 1) >> cnt2)
         ie = e.inv()
@@ -68,7 +68,7 @@ def _butterfly_inv(a: typing.List[Modint]) -> None:
             ies[i - 2] = ie
             e = e * e
             ie = ie * ie
-        sum_ie = [0] * 30
+        sum_ie = [Modint(0)] * 30
         now = Modint(1)
         for i in range(cnt2 - 2):
             sum_ie[i] = ies[i] * now

atcoder/dsu.py

@@ -54,7 +54,7 @@ def size(self, a: int) -> int:
     def groups(self) -> typing.List[typing.List[int]]:
         leader_buf = [self.leader(i) for i in range(self._n)]
 
-        result = [[] for _ in range(self._n)]
+        result: typing.List[typing.List[int]] = [[] for _ in range(self._n)]
         for i in range(self._n):
             result[leader_buf[i]].append(i)
 

atcoder/lazysegtree.py

@@ -82,7 +82,7 @@ def all_prod(self) -> typing.Any:
         return self._d[1]
 
     def apply(self, left: int, right: typing.Optional[int] = None,
-              f: typing.Optional[typing.Any] = None):
+              f: typing.Optional[typing.Any] = None) -> None:
         assert f is not None
 
         if right is None:

atcoder/maxflow.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import NamedTuple, Optional, List
+from typing import NamedTuple, Optional, List, cast
 
 
 class MFGraph:
@@ -38,7 +38,7 @@ def add_edge(self, src: int, dst: int, cap: int) -> int:
     def get_edge(self, i: int) -> Edge:
         assert 0 <= i < len(self._edges)
         e = self._edges[i]
-        re = e.rev
+        re = cast(MFGraph._Edge, e.rev)
         return MFGraph.Edge(
             re.dst,
             e.dst,
@@ -54,14 +54,15 @@ def change_edge(self, i: int, new_cap: int, new_flow: int) -> None:
         assert 0 <= new_flow <= new_cap
         e = self._edges[i]
         e.cap = new_cap - new_flow
+        assert e.rev is not None
         e.rev.cap = new_flow
 
     def flow(self, s: int, t: int, flow_limit: Optional[int] = None) -> int:
         assert 0 <= s < self._n
         assert 0 <= t < self._n
         assert s != t
         if flow_limit is None:
-            flow_limit = sum(e.cap for e in self._g[s])
+            flow_limit = cast(int, sum(e.cap for e in self._g[s]))
 
         current_edge = [0] * self._n
         level = [0] * self._n
@@ -91,20 +92,21 @@ def bfs() -> bool:
 
         def dfs(lim: int) -> int:
             stack = []
-            edge_stack = []
+            edge_stack: List[MFGraph._Edge] = []
             stack.append(t)
             while stack:
                 v = stack[-1]
                 if v == s:
                     flow = min(lim, min(e.cap for e in edge_stack))
                     for e in edge_stack:
                         e.cap -= flow
+                        assert e.rev is not None
                         e.rev.cap += flow
                     return flow
                 next_level = level[v] - 1
                 while current_edge[v] < len(self._g[v]):
                     e = self._g[v][current_edge[v]]
-                    re = e.rev
+                    re = cast(MFGraph._Edge, e.rev)
                     if level[e.dst] != next_level or re.cap == 0:
                         current_edge[v] += 1
                         continue

atcoder/mincostflow.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import NamedTuple, Optional, List
+from typing import NamedTuple, Optional, List, Tuple, cast
 from heapq import heappush, heappop
 
 
@@ -41,7 +41,7 @@ def add_edge(self, src: int, dst: int, cap: int, cost: int) -> int:
     def get_edge(self, i: int) -> Edge:
         assert 0 <= i < len(self._edges)
         e = self._edges[i]
-        re = e.rev
+        re = cast(MCFGraph._Edge, e.rev)
         return MCFGraph.Edge(
             re.dst,
             e.dst,
@@ -53,18 +53,18 @@ def get_edge(self, i: int) -> Edge:
     def edges(self) -> List[Edge]:
         return [self.get_edge(i) for i in range(len(self._edges))]
 
-    def flow(self, s: int, t: int, flow_limit: Optional[int] = None) -> (int, int):
+    def flow(self, s: int, t: int, flow_limit: Optional[int] = None) -> Tuple[int, int]:
         return self.slope(s, t, flow_limit)[-1]
 
-    def slope(self, s: int, t: int, flow_limit: Optional[int] = None) -> List[(int, int)]:
+    def slope(self, s: int, t: int, flow_limit: Optional[int] = None) -> List[Tuple[int, int]]:
         assert 0 <= s < self._n
         assert 0 <= t < self._n
         assert s != t
         if flow_limit is None:
-            flow_limit = sum(e.cap for e in self._g[s])
+            flow_limit = cast(int, sum(e.cap for e in self._g[s]))
 
         dual = [0] * self._n
-        prev: List[Optional[(int, MCFGraph._Edge)]] = [None] * self._n
+        prev: List[Optional[Tuple[int, MCFGraph._Edge]]] = [None] * self._n
 
         def refine_dual() -> bool:
             pq = [(0, s)]
@@ -95,7 +95,7 @@ def refine_dual() -> bool:
             dist_t = dist[t]
             for v in range(self._n):
                 if visited[v]:
-                    dual[v] -= dist_t - dist[v]
+                    dual[v] -= cast(int, dist_t) - cast(int, dist[v])
             return True
 
         flow = 0
@@ -108,13 +108,14 @@ def refine_dual() -> bool:
             f = flow_limit - flow
             v = t
             while prev[v] is not None:
-                u, e = prev[v]
+                u, e = cast(Tuple[int, MCFGraph._Edge], prev[v])
                 f = min(f, e.cap)
                 v = u
             v = t
             while prev[v] is not None:
-                u, e = prev[v]
+                u, e = cast(Tuple[int, MCFGraph._Edge], prev[v])
                 e.cap -= f
+                assert e.rev is not None
                 e.rev.cap += f
                 v = u
             c = -dual[s]

atcoder/modint.py

@@ -5,7 +5,7 @@
 
 
 class ModContext:
-    context = []
+    context: typing.List[int] = []
 
     def __init__(self, mod: int) -> None:
         assert 1 <= mod
@@ -120,13 +120,13 @@ def __floordiv__(self, rhs: typing.Union[Modint, int]) -> Modint:
             inv = atcoder._math._inv_gcd(rhs, self._mod)[1]
         return Modint(self._v * inv)
 
-    def __eq__(self, rhs: typing.Union[Modint, int]) -> bool:
+    def __eq__(self, rhs: typing.Union[Modint, int]) -> bool:  # type: ignore
         if isinstance(rhs, Modint):
             return self._v == rhs._v
         else:
             return self._v == rhs
 
-    def __ne__(self, rhs: typing.Union[Modint, int]) -> bool:
+    def __ne__(self, rhs: typing.Union[Modint, int]) -> bool:  # type: ignore
         if isinstance(rhs, Modint):
             return self._v != rhs._v
         else:

atcoder/string.py

@@ -15,7 +15,7 @@ def _sa_doubling(s: typing.List[int]) -> typing.List[int]:
     tmp = [0] * n
     k = 1
     while k < n:
-        def cmp(x: int, y: int) -> bool:
+        def cmp(x: int, y: int) -> int:
             if rnk[x] != rnk[y]:
                 return rnk[x] - rnk[y]
             rx = rnk[x + k] if x + k < n else -1
@@ -170,7 +170,11 @@ def suffix_array(s: typing.Union[str, typing.List[int]],
     elif upper is None:
         n = len(s)
         idx = list(range(n))
-        idx.sort(key=functools.cmp_to_key(lambda l, r: s[l] - s[r]))
+
+        def cmp(left: int, right: int) -> int:
+            return typing.cast(int, s[left]) - typing.cast(int, s[right])
+
+        idx.sort(key=functools.cmp_to_key(cmp))
         s2 = [0] * n
         now = 0
         for i in range(n):

setup.cfg

@@ -11,4 +11,9 @@ packages = find:
 lint =
     flake8
     pep8-naming
+    mypy
 test = pytest
+
+[mypy]
+disallow_untyped_defs = True
+ignore_missing_imports = True

setup.py

@@ -1,3 +1,3 @@
-import setuptools
+import setuptools  # type: ignore
 
 setuptools.setup()