argmap.py 10.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247
  1. """Utilities for mapping between actual and formal arguments (and their types)."""
  2. from __future__ import annotations
  3. from typing import TYPE_CHECKING, Callable, Sequence
  4. from mypy import nodes
  5. from mypy.maptype import map_instance_to_supertype
  6. from mypy.types import (
  7. AnyType,
  8. Instance,
  9. ParamSpecType,
  10. TupleType,
  11. Type,
  12. TypedDictType,
  13. TypeOfAny,
  14. get_proper_type,
  15. )
  16. if TYPE_CHECKING:
  17. from mypy.infer import ArgumentInferContext
  18. def map_actuals_to_formals(
  19. actual_kinds: list[nodes.ArgKind],
  20. actual_names: Sequence[str | None] | None,
  21. formal_kinds: list[nodes.ArgKind],
  22. formal_names: Sequence[str | None],
  23. actual_arg_type: Callable[[int], Type],
  24. ) -> list[list[int]]:
  25. """Calculate mapping between actual (caller) args and formals.
  26. The result contains a list of caller argument indexes mapping to each
  27. callee argument index, indexed by callee index.
  28. The caller_arg_type argument should evaluate to the type of the actual
  29. argument type with the given index.
  30. """
  31. nformals = len(formal_kinds)
  32. formal_to_actual: list[list[int]] = [[] for i in range(nformals)]
  33. ambiguous_actual_kwargs: list[int] = []
  34. fi = 0
  35. for ai, actual_kind in enumerate(actual_kinds):
  36. if actual_kind == nodes.ARG_POS:
  37. if fi < nformals:
  38. if not formal_kinds[fi].is_star():
  39. formal_to_actual[fi].append(ai)
  40. fi += 1
  41. elif formal_kinds[fi] == nodes.ARG_STAR:
  42. formal_to_actual[fi].append(ai)
  43. elif actual_kind == nodes.ARG_STAR:
  44. # We need to know the actual type to map varargs.
  45. actualt = get_proper_type(actual_arg_type(ai))
  46. if isinstance(actualt, TupleType):
  47. # A tuple actual maps to a fixed number of formals.
  48. for _ in range(len(actualt.items)):
  49. if fi < nformals:
  50. if formal_kinds[fi] != nodes.ARG_STAR2:
  51. formal_to_actual[fi].append(ai)
  52. else:
  53. break
  54. if formal_kinds[fi] != nodes.ARG_STAR:
  55. fi += 1
  56. else:
  57. # Assume that it is an iterable (if it isn't, there will be
  58. # an error later).
  59. while fi < nformals:
  60. if formal_kinds[fi].is_named(star=True):
  61. break
  62. else:
  63. formal_to_actual[fi].append(ai)
  64. if formal_kinds[fi] == nodes.ARG_STAR:
  65. break
  66. fi += 1
  67. elif actual_kind.is_named():
  68. assert actual_names is not None, "Internal error: named kinds without names given"
  69. name = actual_names[ai]
  70. if name in formal_names:
  71. formal_to_actual[formal_names.index(name)].append(ai)
  72. elif nodes.ARG_STAR2 in formal_kinds:
  73. formal_to_actual[formal_kinds.index(nodes.ARG_STAR2)].append(ai)
  74. else:
  75. assert actual_kind == nodes.ARG_STAR2
  76. actualt = get_proper_type(actual_arg_type(ai))
  77. if isinstance(actualt, TypedDictType):
  78. for name in actualt.items:
  79. if name in formal_names:
  80. formal_to_actual[formal_names.index(name)].append(ai)
  81. elif nodes.ARG_STAR2 in formal_kinds:
  82. formal_to_actual[formal_kinds.index(nodes.ARG_STAR2)].append(ai)
  83. else:
  84. # We don't exactly know which **kwargs are provided by the
  85. # caller, so we'll defer until all the other unambiguous
  86. # actuals have been processed
  87. ambiguous_actual_kwargs.append(ai)
  88. if ambiguous_actual_kwargs:
  89. # Assume the ambiguous kwargs will fill the remaining arguments.
  90. #
  91. # TODO: If there are also tuple varargs, we might be missing some potential
  92. # matches if the tuple was short enough to not match everything.
  93. unmatched_formals = [
  94. fi
  95. for fi in range(nformals)
  96. if (
  97. formal_names[fi]
  98. and (
  99. not formal_to_actual[fi]
  100. or actual_kinds[formal_to_actual[fi][0]] == nodes.ARG_STAR
  101. )
  102. and formal_kinds[fi] != nodes.ARG_STAR
  103. )
  104. or formal_kinds[fi] == nodes.ARG_STAR2
  105. ]
  106. for ai in ambiguous_actual_kwargs:
  107. for fi in unmatched_formals:
  108. formal_to_actual[fi].append(ai)
  109. return formal_to_actual
  110. def map_formals_to_actuals(
  111. actual_kinds: list[nodes.ArgKind],
  112. actual_names: Sequence[str | None] | None,
  113. formal_kinds: list[nodes.ArgKind],
  114. formal_names: list[str | None],
  115. actual_arg_type: Callable[[int], Type],
  116. ) -> list[list[int]]:
  117. """Calculate the reverse mapping of map_actuals_to_formals."""
  118. formal_to_actual = map_actuals_to_formals(
  119. actual_kinds, actual_names, formal_kinds, formal_names, actual_arg_type
  120. )
  121. # Now reverse the mapping.
  122. actual_to_formal: list[list[int]] = [[] for _ in actual_kinds]
  123. for formal, actuals in enumerate(formal_to_actual):
  124. for actual in actuals:
  125. actual_to_formal[actual].append(formal)
  126. return actual_to_formal
  127. class ArgTypeExpander:
  128. """Utility class for mapping actual argument types to formal arguments.
  129. One of the main responsibilities is to expand caller tuple *args and TypedDict
  130. **kwargs, and to keep track of which tuple/TypedDict items have already been
  131. consumed.
  132. Example:
  133. def f(x: int, *args: str) -> None: ...
  134. f(*(1, 'x', 1.1))
  135. We'd call expand_actual_type three times:
  136. 1. The first call would provide 'int' as the actual type of 'x' (from '1').
  137. 2. The second call would provide 'str' as one of the actual types for '*args'.
  138. 2. The third call would provide 'float' as one of the actual types for '*args'.
  139. A single instance can process all the arguments for a single call. Each call
  140. needs a separate instance since instances have per-call state.
  141. """
  142. def __init__(self, context: ArgumentInferContext) -> None:
  143. # Next tuple *args index to use.
  144. self.tuple_index = 0
  145. # Keyword arguments in TypedDict **kwargs used.
  146. self.kwargs_used: set[str] = set()
  147. # Type context for `*` and `**` arg kinds.
  148. self.context = context
  149. def expand_actual_type(
  150. self,
  151. actual_type: Type,
  152. actual_kind: nodes.ArgKind,
  153. formal_name: str | None,
  154. formal_kind: nodes.ArgKind,
  155. ) -> Type:
  156. """Return the actual (caller) type(s) of a formal argument with the given kinds.
  157. If the actual argument is a tuple *args, return the next individual tuple item that
  158. maps to the formal arg.
  159. If the actual argument is a TypedDict **kwargs, return the next matching typed dict
  160. value type based on formal argument name and kind.
  161. This is supposed to be called for each formal, in order. Call multiple times per
  162. formal if multiple actuals map to a formal.
  163. """
  164. original_actual = actual_type
  165. actual_type = get_proper_type(actual_type)
  166. if actual_kind == nodes.ARG_STAR:
  167. if isinstance(actual_type, Instance) and actual_type.args:
  168. from mypy.subtypes import is_subtype
  169. if is_subtype(actual_type, self.context.iterable_type):
  170. return map_instance_to_supertype(
  171. actual_type, self.context.iterable_type.type
  172. ).args[0]
  173. else:
  174. # We cannot properly unpack anything other
  175. # than `Iterable` type with `*`.
  176. # Just return `Any`, other parts of code would raise
  177. # a different error for improper use.
  178. return AnyType(TypeOfAny.from_error)
  179. elif isinstance(actual_type, TupleType):
  180. # Get the next tuple item of a tuple *arg.
  181. if self.tuple_index >= len(actual_type.items):
  182. # Exhausted a tuple -- continue to the next *args.
  183. self.tuple_index = 1
  184. else:
  185. self.tuple_index += 1
  186. return actual_type.items[self.tuple_index - 1]
  187. elif isinstance(actual_type, ParamSpecType):
  188. # ParamSpec is valid in *args but it can't be unpacked.
  189. return actual_type
  190. else:
  191. return AnyType(TypeOfAny.from_error)
  192. elif actual_kind == nodes.ARG_STAR2:
  193. from mypy.subtypes import is_subtype
  194. if isinstance(actual_type, TypedDictType):
  195. if formal_kind != nodes.ARG_STAR2 and formal_name in actual_type.items:
  196. # Lookup type based on keyword argument name.
  197. assert formal_name is not None
  198. else:
  199. # Pick an arbitrary item if no specified keyword is expected.
  200. formal_name = (set(actual_type.items.keys()) - self.kwargs_used).pop()
  201. self.kwargs_used.add(formal_name)
  202. return actual_type.items[formal_name]
  203. elif (
  204. isinstance(actual_type, Instance)
  205. and len(actual_type.args) > 1
  206. and is_subtype(actual_type, self.context.mapping_type)
  207. ):
  208. # Only `Mapping` type can be unpacked with `**`.
  209. # Other types will produce an error somewhere else.
  210. return map_instance_to_supertype(actual_type, self.context.mapping_type.type).args[
  211. 1
  212. ]
  213. elif isinstance(actual_type, ParamSpecType):
  214. # ParamSpec is valid in **kwargs but it can't be unpacked.
  215. return actual_type
  216. else:
  217. return AnyType(TypeOfAny.from_error)
  218. else:
  219. # No translation for other kinds -- 1:1 mapping.
  220. return original_actual