ENH: add mypy plugin for more precise ufunc signatures

Author: person142

numpy_ufuncs_plugin.py

@@ -0,0 +1,42 @@
+from mypy.nodes import ARG_POS
+from mypy.plugin import Plugin
+from mypy.types import CallableType
+
+UFUNC_ARITIES = {
+    'sin': 1,
+}
+
+
+def ufunc_call_hook(ctx):
+    ufunc_name = ctx.context.callee.name
+    arity = UFUNC_ARITIES.get(ufunc_name)
+    if arity is None:
+        # No extra information; return the signature unmodified.
+        return ctx.default_signature
+
+    # Strip off the *args and replace it with the correct number of
+    # positional arguments.
+    arg_kinds = [ARG_POS] * arity + ctx.default_signature.arg_kinds[1:]
+    arg_names = (
+        [f'x{i}' for i in range(arity)]
+        + ctx.default_signature.arg_names[1:]
+    )
+    arg_types = (
+        [ctx.default_signature.arg_types[0]] * arity
+        + ctx.default_signature.arg_types[1:]
+    )
+    return ctx.default_signature.copy_modified(
+        arg_kinds=arg_kinds,
+        arg_names=arg_names,
+        arg_types=arg_types,
+    )
+
+
+class UFuncPlugin(Plugin):
+    def get_method_signature_hook(self, method):
+        if method == 'numpy.ufunc.__call__':
+            return ufunc_call_hook
+
+
+def plugin(version):
+    return UFuncPlugin

setup.py

@@ -20,6 +20,7 @@ def find_stubs(package):
     license='BSD',
     version="0.0.1",
     packages=['numpy-stubs'],
+    py_modules=['numpy_ufuncs_plugin'],
     # PEP 561 requires these
     install_requires=['numpy>=1.14.0'],
     package_data=find_stubs('numpy-stubs'),

tests/fail/ufuncs.py

@@ -3,3 +3,4 @@
 np.sin.nin + 'foo'  # E: Unsupported operand types
 np.sin(1, foo='bar')  # E: Unexpected keyword argument
 np.sin(1, extobj=['foo', 'foo', 'foo'])  # E: incompatible type
+np.sin(1, 1)  # E: Too many positional arguments

tests/mypy.ini

@@ -0,0 +1,2 @@
+[mypy]
+plugins = numpy_ufuncs_plugin

tests/test_stubs.py

@@ -9,6 +9,7 @@
 PASS_DIR = os.path.join(TESTS_DIR, "pass")
 FAIL_DIR = os.path.join(TESTS_DIR, "fail")
 REVEAL_DIR = os.path.join(TESTS_DIR, "reveal")
+CONF = os.path.join(TESTS_DIR, "mypy.ini")
 
 
 def get_test_cases(directory):
@@ -38,7 +39,7 @@ def get_test_cases(directory):
 
 @pytest.mark.parametrize("path,py2_arg", get_test_cases(PASS_DIR))
 def test_success(path, py2_arg):
-    stdout, stderr, exitcode = api.run([path] + py2_arg)
+    stdout, stderr, exitcode = api.run([path, "--conf", CONF] + py2_arg)
     assert exitcode == 0, stdout
     assert re.match(
         r'Success: no issues found in \d+ source files?',
@@ -48,7 +49,7 @@ def test_success(path, py2_arg):
 
 @pytest.mark.parametrize("path,py2_arg", get_test_cases(FAIL_DIR))
 def test_fail(path, py2_arg):
-    stdout, stderr, exitcode = api.run([path] + py2_arg)
+    stdout, stderr, exitcode = api.run([path, "--conf", CONF] + py2_arg)
 
     assert exitcode != 0
 
@@ -85,7 +86,7 @@ def test_fail(path, py2_arg):
 
 @pytest.mark.parametrize("path,py2_arg", get_test_cases(REVEAL_DIR))
 def test_reveal(path, py2_arg):
-    stdout, stderr, exitcode = api.run([path] + py2_arg)
+    stdout, stderr, exitcode = api.run([path, "--conf", CONF] + py2_arg)
 
     with open(path) as fin:
         lines = fin.readlines()