infer.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. """Utilities for type argument inference."""
  2. from __future__ import annotations
  3. from typing import NamedTuple, Sequence
  4. from mypy.constraints import (
  5. SUBTYPE_OF,
  6. SUPERTYPE_OF,
  7. infer_constraints,
  8. infer_constraints_for_callable,
  9. )
  10. from mypy.nodes import ArgKind
  11. from mypy.solve import solve_constraints
  12. from mypy.types import CallableType, Instance, Type, TypeVarId
  13. class ArgumentInferContext(NamedTuple):
  14. """Type argument inference context.
  15. We need this because we pass around ``Mapping`` and ``Iterable`` types.
  16. These types are only known by ``TypeChecker`` itself.
  17. It is required for ``*`` and ``**`` argument inference.
  18. https://github.com/python/mypy/issues/11144
  19. """
  20. mapping_type: Instance
  21. iterable_type: Instance
  22. def infer_function_type_arguments(
  23. callee_type: CallableType,
  24. arg_types: Sequence[Type | None],
  25. arg_kinds: list[ArgKind],
  26. formal_to_actual: list[list[int]],
  27. context: ArgumentInferContext,
  28. strict: bool = True,
  29. ) -> list[Type | None]:
  30. """Infer the type arguments of a generic function.
  31. Return an array of lower bound types for the type variables -1 (at
  32. index 0), -2 (at index 1), etc. A lower bound is None if a value
  33. could not be inferred.
  34. Arguments:
  35. callee_type: the target generic function
  36. arg_types: argument types at the call site (each optional; if None,
  37. we are not considering this argument in the current pass)
  38. arg_kinds: nodes.ARG_* values for arg_types
  39. formal_to_actual: mapping from formal to actual variable indices
  40. """
  41. # Infer constraints.
  42. constraints = infer_constraints_for_callable(
  43. callee_type, arg_types, arg_kinds, formal_to_actual, context
  44. )
  45. # Solve constraints.
  46. type_vars = callee_type.type_var_ids()
  47. return solve_constraints(type_vars, constraints, strict)
  48. def infer_type_arguments(
  49. type_var_ids: list[TypeVarId], template: Type, actual: Type, is_supertype: bool = False
  50. ) -> list[Type | None]:
  51. # Like infer_function_type_arguments, but only match a single type
  52. # against a generic type.
  53. constraints = infer_constraints(template, actual, SUPERTYPE_OF if is_supertype else SUBTYPE_OF)
  54. return solve_constraints(type_var_ids, constraints)