Fix the int.__pow__ plugin (python#7632)

msullivan · web-flow · commit 61e08ccea1bc · 2019-10-04T19:56:05.000-07:00
A recent typeshed change added a modulo argument to int.__pow__ which broke the plugin. Tested manually. I didn't add a test because to be useful it would need to be a cmdline test (so it uses the real typeshed) and this doesn't seem worth adding a slow cmdline test for. Happy to add it if someone disagrees though. Fixes python#7621.
diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py
@@ -365,8 +365,10 @@ def typed_dict_update_signature_callback(ctx: MethodSigContext) -> CallableType:
 
 def int_pow_callback(ctx: MethodContext) -> Type:
     """Infer a more precise return type for int.__pow__."""
-    if (len(ctx.arg_types) == 1
-            and len(ctx.arg_types[0]) == 1):
+    # int.__pow__ has an optional modulo argument,
+    # so we expect 2 argument positions
+    if (len(ctx.arg_types) == 2
+            and len(ctx.arg_types[0]) == 1 and len(ctx.arg_types[1]) == 0):
         arg = ctx.args[0][0]
         if isinstance(arg, IntExpr):
             exponent = arg.value
diff --git a/mypy/test/testcmdline.py b/mypy/test/testcmdline.py
@@ -89,6 +89,8 @@ def test_python_cmdline(testcase: DataDrivenTestCase, step: int) -> None:
         if obvious_result != result:
             out.append('== Return code: {}'.format(result))
         expected_out = testcase.output if step == 1 else testcase.output2[step]
+        # Strip "tmp/" out of the test so that # E: works...
+        expected_out = [s.replace("tmp" + os.sep, "") for s in expected_out]
         assert_string_arrays_equal(expected_out, out,
                                    'Invalid output ({}, line {}){}'.format(
                                        testcase.file, testcase.line,
diff --git a/test-data/unit/check-expressions.test b/test-data/unit/check-expressions.test
@@ -2311,25 +2311,6 @@ class B: ...
 
 [builtins fixtures/dict.pyi]
 
--- Type checker default plugin
--- ---------------------------
-
-
-[case testIntPow]
-a = 1
-b = a + 2
-reveal_type(a**0) # N: Revealed type is 'builtins.int'
-reveal_type(a**1) # N: Revealed type is 'builtins.int'
-reveal_type(a**2) # N: Revealed type is 'builtins.int'
-reveal_type(a**-0) # N: Revealed type is 'builtins.int'
-reveal_type(a**-1) # N: Revealed type is 'builtins.float'
-reveal_type(a**(-2)) # N: Revealed type is 'builtins.float'
-reveal_type(a**b) # N: Revealed type is 'Any'
-reveal_type(a.__pow__(2)) # N: Revealed type is 'builtins.int'
-reveal_type(a.__pow__(a)) # N: Revealed type is 'Any'
-a.__pow__() # E: Too few arguments for "__pow__" of "int"
-[builtins fixtures/ops.pyi]
-
 [case testTypeAnnotationNeededMultipleAssignment]
 x, y = [], [] # E: Need type annotation for 'x' (hint: "x: List[<type>] = ...") \
             # E: Need type annotation for 'y' (hint: "y: List[<type>] = ...")
diff --git a/test-data/unit/cmdline.test b/test-data/unit/cmdline.test
@@ -578,6 +578,25 @@ python_version = 3.4
 python_version = 3.6
 [out]
 
+-- This should be a dumping ground for tests of plugins that are sensitive to
+-- typeshed changes.
+[case testTypeshedSensitivePlugins]
+# cmd: mypy int_pow.py
+
+[file int_pow.py]
+a = 1
+b = a + 2
+reveal_type(a**0) # N: Revealed type is 'builtins.int'
+reveal_type(a**1) # N: Revealed type is 'builtins.int'
+reveal_type(a**2) # N: Revealed type is 'builtins.int'
+reveal_type(a**-0) # N: Revealed type is 'builtins.int'
+reveal_type(a**-1) # N: Revealed type is 'builtins.float'
+reveal_type(a**(-2)) # N: Revealed type is 'builtins.float'
+reveal_type(a**b) # N: Revealed type is 'Any'
+reveal_type(a.__pow__(2)) # N: Revealed type is 'builtins.int'
+reveal_type(a.__pow__(a)) # N: Revealed type is 'Any'
+a.__pow__() # E: Too few arguments for "__pow__" of "int"
+
 [case testDisallowAnyUnimported]
 # cmd: mypy main.py
 [file mypy.ini]
diff --git a/test-data/unit/fixtures/ops.pyi b/test-data/unit/fixtures/ops.pyi
@@ -1,4 +1,4 @@
-from typing import overload, Any, Generic, Sequence, Tuple, TypeVar
+from typing import overload, Any, Generic, Sequence, Tuple, TypeVar, Optional
 
 Tco = TypeVar('Tco', covariant=True)
 
@@ -46,7 +46,7 @@ class int:
     def __rtruediv__(self, x: 'int') -> 'int': pass
     def __mod__(self, x: 'int') -> 'int': pass
     def __floordiv__(self, x: 'int') -> 'int': pass
-    def __pow__(self, x: 'int') -> Any: pass
+    def __pow__(self, x: 'int', __modulo: Optional[int] = ...) -> Any: pass
     def __pos__(self) -> 'int': pass
     def __neg__(self) -> 'int': pass
     def __eq__(self, x: object) -> bool: pass
