Generalize the 'open' plugin for 'pathlib.Path.open' (python#7643)

Michael0x2a · gvanrossum · commit 83748a4d546c · 2019-10-09T11:27:52.000-07:00
This pull request adds a plugin to infer a more precise
return type for the `pathlib.Path.open(...)` method.

This method is actually nearly identical to the builtin
`open(...)` method, with the only difference being that
`pathlib.Path.open(...)` method doesn't have the `file`
parameter. So, I refactored the logic in both plugins
into a shared helper method.
diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py
@@ -2,9 +2,10 @@
 from typing import Callable, Optional, List
 
 from mypy import message_registry
-from mypy.nodes import StrExpr, IntExpr, DictExpr, UnaryExpr
+from mypy.nodes import Expression, StrExpr, IntExpr, DictExpr, UnaryExpr
 from mypy.plugin import (
-    Plugin, FunctionContext, MethodContext, MethodSigContext, AttributeContext, ClassDefContext
+    Plugin, FunctionContext, MethodContext, MethodSigContext, AttributeContext, ClassDefContext,
+    CheckerPluginInterface,
 )
 from mypy.plugins.common import try_getting_str_literals
 from mypy.types import (
@@ -66,6 +67,8 @@ def get_method_hook(self, fullname: str
             return ctypes.array_getitem_callback
         elif fullname == 'ctypes.Array.__iter__':
             return ctypes.array_iter_callback
+        elif fullname == 'pathlib.Path.open':
+            return path_open_callback
         return None
 
     def get_attribute_hook(self, fullname: str
@@ -101,23 +104,55 @@ def get_class_decorator_hook(self, fullname: str
 
 
 def open_callback(ctx: FunctionContext) -> Type:
-    """Infer a better return type for 'open'.
-
-    Infer TextIO or BinaryIO as the return value if the mode argument is not
-    given or is a literal.
+    """Infer a better return type for 'open'."""
+    return _analyze_open_signature(
+        arg_types=ctx.arg_types,
+        args=ctx.args,
+        mode_arg_index=1,
+        default_return_type=ctx.default_return_type,
+        api=ctx.api,
+    )
+
+
+def path_open_callback(ctx: MethodContext) -> Type:
+    """Infer a better return type for 'pathlib.Path.open'."""
+    return _analyze_open_signature(
+        arg_types=ctx.arg_types,
+        args=ctx.args,
+        mode_arg_index=0,
+        default_return_type=ctx.default_return_type,
+        api=ctx.api,
+    )
+
+
+def _analyze_open_signature(arg_types: List[List[Type]],
+                            args: List[List[Expression]],
+                            mode_arg_index: int,
+                            default_return_type: Type,
+                            api: CheckerPluginInterface,
+                            ) -> Type:
+    """A helper for analyzing any function that has approximately
+    the same signature as the builtin 'open(...)' function.
+
+    Currently, the only thing the caller can customize is the index
+    of the 'mode' argument. If the mode argument is omitted or is a
+    string literal, we refine the return type to either 'TextIO' or
+    'BinaryIO' as appropriate.
     """
     mode = None
-    if not ctx.arg_types or len(ctx.arg_types[1]) != 1:
+    if not arg_types or len(arg_types[mode_arg_index]) != 1:
         mode = 'r'
-    elif isinstance(ctx.args[1][0], StrExpr):
-        mode = ctx.args[1][0].value
+    else:
+        mode_expr = args[mode_arg_index][0]
+        if isinstance(mode_expr, StrExpr):
+            mode = mode_expr.value
     if mode is not None:
-        assert isinstance(ctx.default_return_type, Instance)  # type: ignore
+        assert isinstance(default_return_type, Instance)  # type: ignore
         if 'b' in mode:
-            return ctx.api.named_generic_type('typing.BinaryIO', [])
+            return api.named_generic_type('typing.BinaryIO', [])
         else:
-            return ctx.api.named_generic_type('typing.TextIO', [])
-    return ctx.default_return_type
+            return api.named_generic_type('typing.TextIO', [])
+    return default_return_type
 
 
 def contextmanager_callback(ctx: FunctionContext) -> Type:
diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test
@@ -299,6 +299,32 @@ _testOpenReturnTypeInferenceSpecialCases.py:2: note: Revealed type is 'typing.Bi
 _testOpenReturnTypeInferenceSpecialCases.py:3: note: Revealed type is 'typing.BinaryIO'
 _testOpenReturnTypeInferenceSpecialCases.py:5: note: Revealed type is 'typing.IO[Any]'
 
+[case testPathOpenReturnTypeInference]
+from pathlib import Path
+p = Path("x")
+reveal_type(p.open())
+reveal_type(p.open('r'))
+reveal_type(p.open('rb'))
+mode = 'rb'
+reveal_type(p.open(mode))
+[out]
+_program.py:3: note: Revealed type is 'typing.TextIO'
+_program.py:4: note: Revealed type is 'typing.TextIO'
+_program.py:5: note: Revealed type is 'typing.BinaryIO'
+_program.py:7: note: Revealed type is 'typing.IO[Any]'
+
+[case testPathOpenReturnTypeInferenceSpecialCases]
+from pathlib import Path
+p = Path("x")
+reveal_type(p.open(mode='rb', errors='replace'))
+reveal_type(p.open(errors='replace', mode='rb'))
+mode = 'rb'
+reveal_type(p.open(mode=mode, errors='replace'))
+[out]
+_program.py:3: note: Revealed type is 'typing.BinaryIO'
+_program.py:4: note: Revealed type is 'typing.BinaryIO'
+_program.py:6: note: Revealed type is 'typing.IO[Any]'
+
 [case testGenericPatterns]
 from typing import Pattern
 import re
