Make black able to compile and run with mypyc

The changes made fall into a couple categories:
 * Fixing actual type mistakes that slip through the cracks
 * Working around a couple mypy bugs (the most annoying of which being
   that we need to add type annotations in a number of places where
   variables are initialized to None

Co-authored-by: Sanjit Kalapatapu <sanjitkal@gmail.com>
Co-authored-by: Michael J. Sullivan <sully@msully.net>

.flake8

@@ -3,3 +3,7 @@ ignore = E203, E266, E501, W503
 max-line-length = 80
 max-complexity = 18
 select = B,C,E,F,W,T4,B9
+# We need to configure the mypy.ini because the flake8-mypy's default
+# options don't properly override it, so if we don't specify it we get
+# half of the config from mypy.ini and half from flake8-mypy.
+mypy_config = mypy.ini

Pipfile

@@ -12,6 +12,8 @@ toml = ">=0.9.4"
 black = {path = ".",extras = ["d"],editable = true}
 aiohttp-cors = "*"
 typed-ast = "==1.4.0"
+typing_extensions = ">=3.7.4"
+mypy_extensions = ">=0.4.3"
 regex = ">=2019.8"
 pathspec = ">=0.6"
 

black.py

@@ -37,6 +37,7 @@
     Union,
     cast,
 )
+from mypy_extensions import mypyc_attr
 
 from appdirs import user_cache_dir
 from attr import dataclass, evolve, Factory
@@ -247,6 +248,17 @@ def read_pyproject_toml(
     return value
 
 
+def target_version_option_callback(
+    c: click.Context, p: Union[click.Option, click.Parameter], v: Tuple[str, ...]
+) -> List[TargetVersion]:
+    """Compute the target versions from a --target-version flag.
+
+    This is its own function because mypy couldn't infer the type correctly
+    when it was a lambda, causing mypyc trouble.
+    """
+    return [TargetVersion[val.upper()] for val in v]
+
+
 @click.command(context_settings=dict(help_option_names=["-h", "--help"]))
 @click.option("-c", "--code", type=str, help="Format the code passed in as a string.")
 @click.option(
@@ -261,7 +273,7 @@ def read_pyproject_toml(
     "-t",
     "--target-version",
     type=click.Choice([v.name.lower() for v in TargetVersion]),
-    callback=lambda c, p, v: [TargetVersion[val.upper()] for val in v],
+    callback=target_version_option_callback,
     multiple=True,
     help=(
         "Python versions that should be supported by Black's output. [default: "
@@ -388,7 +400,7 @@ def main(
     verbose: bool,
     include: str,
     exclude: str,
-    src: Tuple[str],
+    src: Tuple[str, ...],
     config: Optional[str],
 ) -> None:
     """The uncompromising code formatter."""
@@ -470,7 +482,9 @@ def main(
     ctx.exit(report.return_code)
 
 
-def path_empty(src: Tuple[str], quiet: bool, verbose: bool, ctx: click.Context) -> None:
+def path_empty(
+    src: Tuple[str, ...], quiet: bool, verbose: bool, ctx: click.Context
+) -> None:
     """
     Exit if there is no `src` provided for formatting
     """
@@ -639,10 +653,10 @@ def format_file_in_place(
     except NothingChanged:
         return False
 
-    if write_back == write_back.YES:
+    if write_back == WriteBack.YES:
         with open(src, "w", encoding=encoding, newline=newline) as f:
             f.write(dst_contents)
-    elif write_back == write_back.DIFF:
+    elif write_back == WriteBack.DIFF:
         now = datetime.utcnow()
         src_name = f"{src}\t{then} +0000"
         dst_name = f"{src}\t{now} +0000"
@@ -1728,13 +1742,13 @@ def visit_default(self, node: LN) -> Iterator[Line]:
                 self.current_line.append(node)
         yield from super().visit_default(node)
 
-    def visit_INDENT(self, node: Node) -> Iterator[Line]:
+    def visit_INDENT(self, node: Leaf) -> Iterator[Line]:
         """Increase indentation level, maybe yield a line."""
         # In blib2to3 INDENT never holds comments.
         yield from self.line(+1)
         yield from self.visit_default(node)
 
-    def visit_DEDENT(self, node: Node) -> Iterator[Line]:
+    def visit_DEDENT(self, node: Leaf) -> Iterator[Line]:
         """Decrease indentation level, maybe yield a line."""
         # The current line might still wait for trailing comments.  At DEDENT time
         # there won't be any (they would be prefixes on the preceding NEWLINE).
@@ -2462,7 +2476,7 @@ def left_hand_split(line: Line, features: Collection[Feature] = ()) -> Iterator[
     body_leaves: List[Leaf] = []
     head_leaves: List[Leaf] = []
     current_leaves = head_leaves
-    matching_bracket = None
+    matching_bracket: Optional[Leaf] = None
     for leaf in line.leaves:
         if (
             current_leaves is body_leaves
@@ -2505,8 +2519,8 @@ def right_hand_split(
     body_leaves: List[Leaf] = []
     head_leaves: List[Leaf] = []
     current_leaves = tail_leaves
-    opening_bracket = None
-    closing_bracket = None
+    opening_bracket: Optional[Leaf] = None
+    closing_bracket: Optional[Leaf] = None
     for leaf in reversed(line.leaves):
         if current_leaves is body_leaves:
             if leaf is opening_bracket:
@@ -3027,7 +3041,7 @@ def convert_one_fmt_off_pair(node: Node) -> bool:
                     # That happens when one of the `ignored_nodes` ended with a NEWLINE
                     # leaf (possibly followed by a DEDENT).
                     hidden_value = hidden_value[:-1]
-                first_idx = None
+                first_idx: Optional[int] = None
                 for ignored in ignored_nodes:
                     index = ignored.remove()
                     if first_idx is None:
@@ -3398,8 +3412,8 @@ def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[Leaf
     yield omit
 
     length = 4 * line.depth
-    opening_bracket = None
-    closing_bracket = None
+    opening_bracket: Optional[Leaf] = None
+    closing_bracket: Optional[Leaf] = None
     inner_brackets: Set[LeafID] = set()
     for index, leaf, leaf_length in enumerate_with_length(line, reversed=True):
         length += leaf_length
@@ -3796,6 +3810,7 @@ def assert_stable(src: str, dst: str, mode: FileMode) -> None:
         ) from None
 
 
+@mypyc_attr(patchable=True)
 def dump_to_file(*output: str) -> str:
     """Dump `output` to a temporary file. Return path to the file."""
     with tempfile.NamedTemporaryFile(

setup.py

@@ -1,6 +1,7 @@
 # Copyright (C) 2018 Łukasz Langa
 from setuptools import setup
 import sys
+import os
 
 assert sys.version_info >= (3, 6, 0), "black requires Python 3.6+"
 from pathlib import Path  # noqa E402
@@ -15,6 +16,33 @@ def get_long_description() -> str:
         return ld_file.read()
 
 
+USE_MYPYC = False
+# To compile with mypyc, a mypyc checkout must be present on the PYTHONPATH
+if len(sys.argv) > 1 and sys.argv[1] == "--use-mypyc":
+    sys.argv.pop(1)
+    USE_MYPYC = True
+if os.getenv("BLACK_USE_MYPYC", None) == "1":
+    USE_MYPYC = True
+
+if USE_MYPYC:
+    mypyc_targets = [
+        "black.py",
+        "blib2to3/pytree.py",
+        "blib2to3/pygram.py",
+        "blib2to3/pgen2/parse.py",
+        "blib2to3/pgen2/grammar.py",
+        "blib2to3/pgen2/token.py",
+        "blib2to3/pgen2/driver.py",
+        "blib2to3/pgen2/pgen.py",
+    ]
+
+    from mypyc.build import mypycify
+
+    opt_level = os.getenv("MYPYC_OPT_LEVEL", "3")
+    ext_modules = mypycify(mypyc_targets, opt_level=opt_level)
+else:
+    ext_modules = []
+
 setup(
     name="black",
     use_scm_version={
@@ -30,6 +58,7 @@ def get_long_description() -> str:
     url="https://github.com/psf/black",
     license="MIT",
     py_modules=["black", "blackd", "_black_version"],
+    ext_modules=ext_modules,
     packages=["blib2to3", "blib2to3.pgen2"],
     package_data={"blib2to3": ["*.txt"]},
     python_requires=">=3.6",
@@ -42,6 +71,8 @@ def get_long_description() -> str:
         "typed-ast>=1.4.0",
         "regex",
         "pathspec>=0.6, <1",
+        "typing_extensions>=3.7.4",
+        "mypy_extensions>=0.4.3",
     ],
     extras_require={"d": ["aiohttp>=3.3.2", "aiohttp-cors"]},
     test_suite="tests.test_black",

tests/test_black.py

@@ -1540,6 +1540,7 @@ def test_symlink_out_of_root_directory(self) -> None:
         # outside of the `root` directory.
         path.iterdir.return_value = [child]
         child.resolve.return_value = Path("/a/b/c")
+        child.as_posix.return_value = "/a/b/c"
         child.is_symlink.return_value = True
         try:
             list(