127 lines
4.6 KiB
Python
127 lines
4.6 KiB
Python
|
from typing import Optional, Dict, Union
|
||
|
|
from mypy.types import (
|
||
|
|
TypeVarLikeType, TypeVarType, ParamSpecType, ParamSpecFlavor, TypeVarId, TypeVarTupleType,
|
||
|
|
)
|
||
|
|
from mypy.nodes import (
|
||
|
|
ParamSpecExpr, TypeVarExpr, TypeVarLikeExpr, SymbolTableNode, TypeVarTupleExpr,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class TypeVarLikeScope:
|
||
|
|
"""Scope that holds bindings for type variables and parameter specifications.
|
||
|
|
|
||
|
|
Node fullname -> TypeVarLikeType.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self,
|
||
|
|
parent: 'Optional[TypeVarLikeScope]' = None,
|
||
|
|
is_class_scope: bool = False,
|
||
|
|
prohibited: 'Optional[TypeVarLikeScope]' = None,
|
||
|
|
namespace: str = '') -> None:
|
||
|
|
"""Initializer for TypeVarLikeScope
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
parent: the outer scope for this scope
|
||
|
|
is_class_scope: True if this represents a generic class
|
||
|
|
prohibited: Type variables that aren't strictly in scope exactly,
|
||
|
|
but can't be bound because they're part of an outer class's scope.
|
||
|
|
"""
|
||
|
|
self.scope: Dict[str, TypeVarLikeType] = {}
|
||
|
|
self.parent = parent
|
||
|
|
self.func_id = 0
|
||
|
|
self.class_id = 0
|
||
|
|
self.is_class_scope = is_class_scope
|
||
|
|
self.prohibited = prohibited
|
||
|
|
self.namespace = namespace
|
||
|
|
if parent is not None:
|
||
|
|
self.func_id = parent.func_id
|
||
|
|
self.class_id = parent.class_id
|
||
|
|
|
||
|
|
def get_function_scope(self) -> 'Optional[TypeVarLikeScope]':
|
||
|
|
"""Get the nearest parent that's a function scope, not a class scope"""
|
||
|
|
it: Optional[TypeVarLikeScope] = self
|
||
|
|
while it is not None and it.is_class_scope:
|
||
|
|
it = it.parent
|
||
|
|
return it
|
||
|
|
|
||
|
|
def allow_binding(self, fullname: str) -> bool:
|
||
|
|
if fullname in self.scope:
|
||
|
|
return False
|
||
|
|
elif self.parent and not self.parent.allow_binding(fullname):
|
||
|
|
return False
|
||
|
|
elif self.prohibited and not self.prohibited.allow_binding(fullname):
|
||
|
|
return False
|
||
|
|
return True
|
||
|
|
|
||
|
|
def method_frame(self) -> 'TypeVarLikeScope':
|
||
|
|
"""A new scope frame for binding a method"""
|
||
|
|
return TypeVarLikeScope(self, False, None)
|
||
|
|
|
||
|
|
def class_frame(self, namespace: str) -> 'TypeVarLikeScope':
|
||
|
|
"""A new scope frame for binding a class. Prohibits *this* class's tvars"""
|
||
|
|
return TypeVarLikeScope(self.get_function_scope(), True, self, namespace=namespace)
|
||
|
|
|
||
|
|
def bind_new(self, name: str, tvar_expr: TypeVarLikeExpr) -> TypeVarLikeType:
|
||
|
|
if self.is_class_scope:
|
||
|
|
self.class_id += 1
|
||
|
|
i = self.class_id
|
||
|
|
namespace = self.namespace
|
||
|
|
else:
|
||
|
|
self.func_id -= 1
|
||
|
|
i = self.func_id
|
||
|
|
# TODO: Consider also using namespaces for functions
|
||
|
|
namespace = ''
|
||
|
|
if isinstance(tvar_expr, TypeVarExpr):
|
||
|
|
tvar_def: TypeVarLikeType = TypeVarType(
|
||
|
|
name,
|
||
|
|
tvar_expr.fullname,
|
||
|
|
TypeVarId(i, namespace=namespace),
|
||
|
|
values=tvar_expr.values,
|
||
|
|
upper_bound=tvar_expr.upper_bound,
|
||
|
|
variance=tvar_expr.variance,
|
||
|
|
line=tvar_expr.line,
|
||
|
|
column=tvar_expr.column
|
||
|
|
)
|
||
|
|
elif isinstance(tvar_expr, ParamSpecExpr):
|
||
|
|
tvar_def = ParamSpecType(
|
||
|
|
name,
|
||
|
|
tvar_expr.fullname,
|
||
|
|
i,
|
||
|
|
flavor=ParamSpecFlavor.BARE,
|
||
|
|
upper_bound=tvar_expr.upper_bound,
|
||
|
|
line=tvar_expr.line,
|
||
|
|
column=tvar_expr.column
|
||
|
|
)
|
||
|
|
elif isinstance(tvar_expr, TypeVarTupleExpr):
|
||
|
|
tvar_def = TypeVarTupleType(
|
||
|
|
name,
|
||
|
|
tvar_expr.fullname,
|
||
|
|
i,
|
||
|
|
upper_bound=tvar_expr.upper_bound,
|
||
|
|
line=tvar_expr.line,
|
||
|
|
column=tvar_expr.column
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
assert False
|
||
|
|
self.scope[tvar_expr.fullname] = tvar_def
|
||
|
|
return tvar_def
|
||
|
|
|
||
|
|
def bind_existing(self, tvar_def: TypeVarLikeType) -> None:
|
||
|
|
self.scope[tvar_def.fullname] = tvar_def
|
||
|
|
|
||
|
|
def get_binding(self, item: Union[str, SymbolTableNode]) -> Optional[TypeVarLikeType]:
|
||
|
|
fullname = item.fullname if isinstance(item, SymbolTableNode) else item
|
||
|
|
assert fullname is not None
|
||
|
|
if fullname in self.scope:
|
||
|
|
return self.scope[fullname]
|
||
|
|
elif self.parent is not None:
|
||
|
|
return self.parent.get_binding(fullname)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def __str__(self) -> str:
|
||
|
|
me = ", ".join(f'{k}: {v.name}`{v.id}' for k, v in self.scope.items())
|
||
|
|
if self.parent is None:
|
||
|
|
return me
|
||
|
|
return f"{self.parent} <- {me}"
|