Infer ufunc arity from the ufunc object itself

Author: person142

numpy_ufuncs_plugin.py

@@ -1,28 +1,25 @@
 from mypy.nodes import ARG_POS
 from mypy.plugin import Plugin
 from mypy.types import CallableType
-
-UFUNC_ARITIES = {
-    'sin': 1,
-}
+import numpy as np
 
 
 def ufunc_call_hook(ctx):
     ufunc_name = ctx.context.callee.name
-    arity = UFUNC_ARITIES.get(ufunc_name)
-    if arity is None:
+    ufunc = getattr(np, ufunc_name, None)
+    if ufunc 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_kinds = [ARG_POS] * ufunc.nin + ctx.default_signature.arg_kinds[1:]
     arg_names = (
-        [f'x{i}' for i in range(arity)] +
+        [f'x{i}' for i in range(ufunc.nin)] +
         ctx.default_signature.arg_names[1:]
     )
     arg_types = (
-        [ctx.default_signature.arg_types[0]] * arity +
+        [ctx.default_signature.arg_types[0]] * ufunc.nin +
         ctx.default_signature.arg_types[1:]
     )
     return ctx.default_signature.copy_modified(

tests/mypy.ini

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