From a7069144307ea3c4d17a5d846be11bd8a8901843 Mon Sep 17 00:00:00 2001 From: Stanislav Terliakov <50529348+sterliakov@users.noreply.github.com> Date: Sun, 27 Oct 2024 00:04:35 +0200 Subject: [PATCH] Support ParamSpec mapping with functools.partial (#17355) Follow-up for #17323, resolving a false positive discovered there. Closes #17960. This enables use of `functools.partial` to bind some `*args` or `**kwargs` on a callable typed with `ParamSpec`. --------- Co-authored-by: Shantanu Jain --- mypy/checkexpr.py | 6 +- mypy/plugins/functools.py | 47 ++++++- mypy/types.py | 4 +- .../unit/check-parameter-specification.test | 124 ++++++++++++++++++ 4 files changed, 174 insertions(+), 7 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index d63cf6e782c7..577576a4e5f8 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2377,7 +2377,11 @@ def check_argument_count( # Positional argument when expecting a keyword argument. self.msg.too_many_positional_arguments(callee, context) ok = False - elif callee.param_spec() is not None and not formal_to_actual[i]: + elif ( + callee.param_spec() is not None + and not formal_to_actual[i] + and callee.special_sig != "partial" + ): self.msg.too_few_arguments(callee, context, actual_names) ok = False return ok diff --git a/mypy/plugins/functools.py b/mypy/plugins/functools.py index 4dfeb752b5d2..6a063174bfcb 100644 --- a/mypy/plugins/functools.py +++ b/mypy/plugins/functools.py @@ -8,7 +8,7 @@ import mypy.plugin import mypy.semanal from mypy.argmap import map_actuals_to_formals -from mypy.nodes import ARG_POS, ARG_STAR2, ArgKind, Argument, CallExpr, FuncItem, Var +from mypy.nodes import ARG_POS, ARG_STAR2, ArgKind, Argument, CallExpr, FuncItem, NameExpr, Var from mypy.plugins.common import add_method_to_class from mypy.typeops import get_all_type_vars from mypy.types import ( @@ -16,6 +16,8 @@ CallableType, Instance, Overloaded, + ParamSpecFlavor, + ParamSpecType, Type, TypeOfAny, TypeVarType, @@ -202,6 +204,7 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) - continue can_infer_ids.update({tv.id for tv in get_all_type_vars(arg_type)}) + # special_sig="partial" allows omission of args/kwargs typed with ParamSpec defaulted = fn_type.copy_modified( arg_kinds=[ ( @@ -218,6 +221,7 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) - # Keep TypeVarTuple/ParamSpec to avoid spurious errors on empty args. if tv.id in can_infer_ids or not isinstance(tv, TypeVarType) ], + special_sig="partial", ) if defaulted.line < 0: # Make up a line number if we don't have one @@ -296,10 +300,19 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) - arg_kinds=partial_kinds, arg_names=partial_names, ret_type=ret_type, + special_sig="partial", ) ret = ctx.api.named_generic_type(PARTIAL, [ret_type]) ret = ret.copy_with_extra_attr("__mypy_partial", partially_applied) + if partially_applied.param_spec(): + assert ret.extra_attrs is not None # copy_with_extra_attr above ensures this + attrs = ret.extra_attrs.copy() + if ArgKind.ARG_STAR in actual_arg_kinds: + attrs.immutable.add("__mypy_partial_paramspec_args_bound") + if ArgKind.ARG_STAR2 in actual_arg_kinds: + attrs.immutable.add("__mypy_partial_paramspec_kwargs_bound") + ret.extra_attrs = attrs return ret @@ -314,7 +327,8 @@ def partial_call_callback(ctx: mypy.plugin.MethodContext) -> Type: ): return ctx.default_return_type - partial_type = ctx.type.extra_attrs.attrs["__mypy_partial"] + extra_attrs = ctx.type.extra_attrs + partial_type = get_proper_type(extra_attrs.attrs["__mypy_partial"]) if len(ctx.arg_types) != 2: # *args, **kwargs return ctx.default_return_type @@ -332,11 +346,36 @@ def partial_call_callback(ctx: mypy.plugin.MethodContext) -> Type: actual_arg_kinds.append(ctx.arg_kinds[i][j]) actual_arg_names.append(ctx.arg_names[i][j]) - result = ctx.api.expr_checker.check_call( + result, _ = ctx.api.expr_checker.check_call( callee=partial_type, args=actual_args, arg_kinds=actual_arg_kinds, arg_names=actual_arg_names, context=ctx.context, ) - return result[0] + if not isinstance(partial_type, CallableType) or partial_type.param_spec() is None: + return result + + args_bound = "__mypy_partial_paramspec_args_bound" in extra_attrs.immutable + kwargs_bound = "__mypy_partial_paramspec_kwargs_bound" in extra_attrs.immutable + + passed_paramspec_parts = [ + arg.node.type + for arg in actual_args + if isinstance(arg, NameExpr) + and isinstance(arg.node, Var) + and isinstance(arg.node.type, ParamSpecType) + ] + # ensure *args: P.args + args_passed = any(part.flavor == ParamSpecFlavor.ARGS for part in passed_paramspec_parts) + if not args_bound and not args_passed: + ctx.api.expr_checker.msg.too_few_arguments(partial_type, ctx.context, actual_arg_names) + elif args_bound and args_passed: + ctx.api.expr_checker.msg.too_many_arguments(partial_type, ctx.context) + + # ensure **kwargs: P.kwargs + kwargs_passed = any(part.flavor == ParamSpecFlavor.KWARGS for part in passed_paramspec_parts) + if not kwargs_bound and not kwargs_passed: + ctx.api.expr_checker.msg.too_few_arguments(partial_type, ctx.context, actual_arg_names) + + return result diff --git a/mypy/types.py b/mypy/types.py index 897e19d6ee19..0b010ca9d593 100644 --- a/mypy/types.py +++ b/mypy/types.py @@ -1827,8 +1827,8 @@ class CallableType(FunctionLike): "implicit", # Was this type implicitly generated instead of explicitly # specified by the user? "special_sig", # Non-None for signatures that require special handling - # (currently only value is 'dict' for a signature similar to - # 'dict') + # (currently only values are 'dict' for a signature similar to + # 'dict' and 'partial' for a `functools.partial` evaluation) "from_type_type", # Was this callable generated by analyzing Type[...] # instantiation? "bound_args", # Bound type args, mostly unused but may be useful for diff --git a/test-data/unit/check-parameter-specification.test b/test-data/unit/check-parameter-specification.test index f499bac45102..674e3894940b 100644 --- a/test-data/unit/check-parameter-specification.test +++ b/test-data/unit/check-parameter-specification.test @@ -2338,3 +2338,127 @@ reveal_type(handle_reversed(Child())) # N: Revealed type is "builtins.str" reveal_type(handle_reversed(NotChild())) # N: Revealed type is "builtins.str" [builtins fixtures/paramspec.pyi] + +[case testBindPartial] +from functools import partial +from typing_extensions import ParamSpec +from typing import Callable, TypeVar + +P = ParamSpec("P") +T = TypeVar("T") + +def run(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, **kwargs) + return func2(*args) + +def run2(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, *args) + return func2(**kwargs) + +def run3(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, *args, **kwargs) + return func2() + +def run4(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, *args, **kwargs) + return func2(**kwargs) + +def run_bad(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, *args, **kwargs) + return func2(*args) # E: Too many arguments + +def run_bad2(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, **kwargs) + return func2(**kwargs) # E: Too few arguments + +def run_bad3(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, *args) + return func2() # E: Too few arguments + +[builtins fixtures/paramspec.pyi] + +[case testBindPartialConcatenate] +from functools import partial +from typing_extensions import Concatenate, ParamSpec +from typing import Callable, TypeVar + +P = ParamSpec("P") +T = TypeVar("T") + +def run(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1, **kwargs) + return func2(*args) + +def run2(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, **kwargs) + p = [""] + func2(1, *p) # E: Too few arguments \ + # E: Argument 2 has incompatible type "*List[str]"; expected "P.args" + func2(1, 2, *p) # E: Too few arguments \ + # E: Argument 2 has incompatible type "int"; expected "P.args" \ + # E: Argument 3 has incompatible type "*List[str]"; expected "P.args" + func2(1, *args, *p) # E: Argument 3 has incompatible type "*List[str]"; expected "P.args" + func2(1, *p, *args) # E: Argument 2 has incompatible type "*List[str]"; expected "P.args" + return func2(1, *args) + +def run3(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1, *args) + d = {"":""} + func2(**d) # E: Too few arguments \ + # E: Argument 1 has incompatible type "**Dict[str, str]"; expected "P.kwargs" + return func2(**kwargs) + +def run4(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1) + return func2(*args, **kwargs) + +def run5(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1, *args, **kwargs) + func2() + return func2(**kwargs) + +def run_bad(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, *args) # E: Argument 1 has incompatible type "*P.args"; expected "int" + return func2(1, **kwargs) # E: Argument 1 has incompatible type "int"; expected "P.args" + +def run_bad2(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1, *args) + func2() # E: Too few arguments + func2(*args, **kwargs) # E: Too many arguments + return func2(1, **kwargs) # E: Argument 1 has incompatible type "int"; expected "P.args" + +def run_bad3(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1, **kwargs) + func2() # E: Too few arguments + return func2(1, *args) # E: Argument 1 has incompatible type "int"; expected "P.args" + +def run_bad4(func: Callable[Concatenate[int, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, 1) + func2() # E: Too few arguments + func2(*args) # E: Too few arguments + func2(1, *args) # E: Too few arguments \ + # E: Argument 1 has incompatible type "int"; expected "P.args" + func2(1, **kwargs) # E: Too few arguments \ + # E: Argument 1 has incompatible type "int"; expected "P.args" + return func2(**kwargs) # E: Too few arguments + +[builtins fixtures/paramspec.pyi] + +[case testOtherVarArgs] +from functools import partial +from typing_extensions import Concatenate, ParamSpec +from typing import Callable, TypeVar, Tuple + +P = ParamSpec("P") +T = TypeVar("T") + +def run(func: Callable[Concatenate[int, str, P], T], *args: P.args, **kwargs: P.kwargs) -> T: + func2 = partial(func, **kwargs) + args_prefix: Tuple[int, str] = (1, 'a') + func2(*args_prefix) # E: Too few arguments + func2(*args, *args_prefix) # E: Argument 1 has incompatible type "*P.args"; expected "int" \ + # E: Argument 1 has incompatible type "*P.args"; expected "str" \ + # E: Argument 2 has incompatible type "*Tuple[int, str]"; expected "P.args" + return func2(*args_prefix, *args) + +[builtins fixtures/paramspec.pyi]