Abuse literals to allow inferring ufunc arity from Generic

Author: person142

numpy-stubs/__init__.pyi

@@ -28,6 +28,7 @@ from typing import (
     TypeVar,
     Union,
 )
+from typing_extensions import Literal
 
 if sys.version_info[0] < 3:
     class SupportsBytes: ...
@@ -619,7 +620,10 @@ WRAP: int
 little_endian: int
 tracemalloc_domain: int
 
-class ufunc:
+_Nin = TypeVar('_Nin', bound=int)
+_Nout = TypeVar('_Nout', bound=int)
+
+class ufunc(Generic[_Nin], Generic[_Nout]):
     @property
     def __name__(self) -> str: ...
     def __call__(
@@ -765,7 +769,7 @@ right_shift: ufunc
 rint: ufunc
 sign: ufunc
 signbit: ufunc
-sin: ufunc
+sin: ufunc[Literal[1], Literal[1]]
 sinh: ufunc
 spacing: ufunc
 sqrt: ufunc

numpy_ufuncs_plugin.py

@@ -1,25 +1,30 @@
 from mypy.nodes import ARG_POS
 from mypy.plugin import Plugin
 from mypy.types import CallableType
-import numpy as np
 
 
 def ufunc_call_hook(ctx):
     ufunc_name = ctx.context.callee.name
-    ufunc = getattr(np, ufunc_name, None)
-    if ufunc is None:
-        # No extra information; return the signature unmodified.
+
+    type_info = ctx.type.serialize()
+    nin_arg, nout_arg = type_info['args']
+    if nin_arg['.class'] != 'LiteralType':
+        return ctx.default_signature
+    if nout_arg['.class'] != 'LiteralType':
         return ctx.default_signature
 
+    nin = nin_arg['value']
+    nout = nout_arg['value']
+
     # Strip off the *args and replace it with the correct number of
     # positional arguments.
-    arg_kinds = [ARG_POS] * ufunc.nin + ctx.default_signature.arg_kinds[1:]
+    arg_kinds = [ARG_POS] * nin + ctx.default_signature.arg_kinds[1:]
     arg_names = (
-        [f'x{i}' for i in range(ufunc.nin)] +
+        [f'x{i}' for i in range(nin)] +
         ctx.default_signature.arg_names[1:]
     )
     arg_types = (
-        [ctx.default_signature.arg_types[0]] * ufunc.nin +
+        [ctx.default_signature.arg_types[0]] * nin +
         ctx.default_signature.arg_types[1:]
     )
     return ctx.default_signature.copy_modified(