tvar_scope.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. from __future__ import annotations
  2. from mypy.nodes import (
  3. ParamSpecExpr,
  4. SymbolTableNode,
  5. TypeVarExpr,
  6. TypeVarLikeExpr,
  7. TypeVarTupleExpr,
  8. )
  9. from mypy.types import (
  10. ParamSpecFlavor,
  11. ParamSpecType,
  12. TypeVarId,
  13. TypeVarLikeType,
  14. TypeVarTupleType,
  15. TypeVarType,
  16. )
  17. class TypeVarLikeScope:
  18. """Scope that holds bindings for type variables and parameter specifications.
  19. Node fullname -> TypeVarLikeType.
  20. """
  21. def __init__(
  22. self,
  23. parent: TypeVarLikeScope | None = None,
  24. is_class_scope: bool = False,
  25. prohibited: TypeVarLikeScope | None = None,
  26. namespace: str = "",
  27. ) -> None:
  28. """Initializer for TypeVarLikeScope
  29. Parameters:
  30. parent: the outer scope for this scope
  31. is_class_scope: True if this represents a generic class
  32. prohibited: Type variables that aren't strictly in scope exactly,
  33. but can't be bound because they're part of an outer class's scope.
  34. """
  35. self.scope: dict[str, TypeVarLikeType] = {}
  36. self.parent = parent
  37. self.func_id = 0
  38. self.class_id = 0
  39. self.is_class_scope = is_class_scope
  40. self.prohibited = prohibited
  41. self.namespace = namespace
  42. if parent is not None:
  43. self.func_id = parent.func_id
  44. self.class_id = parent.class_id
  45. def get_function_scope(self) -> TypeVarLikeScope | None:
  46. """Get the nearest parent that's a function scope, not a class scope"""
  47. it: TypeVarLikeScope | None = self
  48. while it is not None and it.is_class_scope:
  49. it = it.parent
  50. return it
  51. def allow_binding(self, fullname: str) -> bool:
  52. if fullname in self.scope:
  53. return False
  54. elif self.parent and not self.parent.allow_binding(fullname):
  55. return False
  56. elif self.prohibited and not self.prohibited.allow_binding(fullname):
  57. return False
  58. return True
  59. def method_frame(self) -> TypeVarLikeScope:
  60. """A new scope frame for binding a method"""
  61. return TypeVarLikeScope(self, False, None)
  62. def class_frame(self, namespace: str) -> TypeVarLikeScope:
  63. """A new scope frame for binding a class. Prohibits *this* class's tvars"""
  64. return TypeVarLikeScope(self.get_function_scope(), True, self, namespace=namespace)
  65. def new_unique_func_id(self) -> int:
  66. """Used by plugin-like code that needs to make synthetic generic functions."""
  67. self.func_id -= 1
  68. return self.func_id
  69. def bind_new(self, name: str, tvar_expr: TypeVarLikeExpr) -> TypeVarLikeType:
  70. if self.is_class_scope:
  71. self.class_id += 1
  72. i = self.class_id
  73. namespace = self.namespace
  74. else:
  75. self.func_id -= 1
  76. i = self.func_id
  77. # TODO: Consider also using namespaces for functions
  78. namespace = ""
  79. if isinstance(tvar_expr, TypeVarExpr):
  80. tvar_def: TypeVarLikeType = TypeVarType(
  81. name=name,
  82. fullname=tvar_expr.fullname,
  83. id=TypeVarId(i, namespace=namespace),
  84. values=tvar_expr.values,
  85. upper_bound=tvar_expr.upper_bound,
  86. default=tvar_expr.default,
  87. variance=tvar_expr.variance,
  88. line=tvar_expr.line,
  89. column=tvar_expr.column,
  90. )
  91. elif isinstance(tvar_expr, ParamSpecExpr):
  92. tvar_def = ParamSpecType(
  93. name,
  94. tvar_expr.fullname,
  95. i,
  96. flavor=ParamSpecFlavor.BARE,
  97. upper_bound=tvar_expr.upper_bound,
  98. default=tvar_expr.default,
  99. line=tvar_expr.line,
  100. column=tvar_expr.column,
  101. )
  102. elif isinstance(tvar_expr, TypeVarTupleExpr):
  103. tvar_def = TypeVarTupleType(
  104. name,
  105. tvar_expr.fullname,
  106. i,
  107. upper_bound=tvar_expr.upper_bound,
  108. tuple_fallback=tvar_expr.tuple_fallback,
  109. default=tvar_expr.default,
  110. line=tvar_expr.line,
  111. column=tvar_expr.column,
  112. )
  113. else:
  114. assert False
  115. self.scope[tvar_expr.fullname] = tvar_def
  116. return tvar_def
  117. def bind_existing(self, tvar_def: TypeVarLikeType) -> None:
  118. self.scope[tvar_def.fullname] = tvar_def
  119. def get_binding(self, item: str | SymbolTableNode) -> TypeVarLikeType | None:
  120. fullname = item.fullname if isinstance(item, SymbolTableNode) else item
  121. assert fullname
  122. if fullname in self.scope:
  123. return self.scope[fullname]
  124. elif self.parent is not None:
  125. return self.parent.get_binding(fullname)
  126. else:
  127. return None
  128. def __str__(self) -> str:
  129. me = ", ".join(f"{k}: {v.name}`{v.id}" for k, v in self.scope.items())
  130. if self.parent is None:
  131. return me
  132. return f"{self.parent} <- {me}"