Class PolyTranslator

mypy/applytype.py:206–303

Make free type variables generic in the type if possible. See docstring for apply_poly() for details.

class PolyTranslator(TypeTranslator):
    """Make free type variables generic in the type if possible.

    See docstring for apply_poly() for details.
    """

    def __init__(
        self,
        poly_tvars: Iterable[TypeVarLikeType],
        bound_tvars: frozenset[TypeVarLikeType] = frozenset(),
        seen_aliases: frozenset[TypeInfo] = frozenset(),
    ) -> None:
        super().__init__()
        self.poly_tvars = set(poly_tvars)
        # This is a simplified version of TypeVarScope used during semantic analysis.
        self.bound_tvars = bound_tvars
        self.seen_aliases = seen_aliases

    def collect_vars(self, t: CallableType | Parameters) -> list[TypeVarLikeType]:
        found_vars = []
        for arg in t.arg_types:
            for tv in get_all_type_vars(arg):
                if isinstance(tv, ParamSpecType):
                    normalized: TypeVarLikeType = tv.copy_modified(
                        flavor=ParamSpecFlavor.BARE, prefix=Parameters([], [], [])
                    )
                else:
                    normalized = tv
                if normalized in self.poly_tvars and normalized not in self.bound_tvars:
                    found_vars.append(normalized)
        return remove_dups(found_vars)

    def visit_callable_type(self, t: CallableType) -> Type:
        found_vars = self.collect_vars(t)
        self.bound_tvars |= set(found_vars)
        result = super().visit_callable_type(t)
        self.bound_tvars -= set(found_vars)

        assert isinstance(result, ProperType) and isinstance(result, CallableType)
        result.variables = result.variables + tuple(found_vars)
        return result

    def visit_type_var(self, t: TypeVarType) -> Type:
        if t in self.poly_tvars and t not in self.bound_tvars:
            raise PolyTranslationError()
        return super().visit_type_var(t)

    def visit_param_spec(self, t: ParamSpecType) -> Type:
        if t in self.poly_tvars and t not in self.bound_tvars:
            raise PolyTranslationError()
        return super().visit_param_spec(t)

    def visit_type_var_tuple(self, t: TypeVarTupleType) -> Type:
        if t in self.poly_tvars and t not in self.bound_tvars:
            raise PolyTranslationError()
        return super().visit_type_var_tuple(t)

    def visit_type_alias_type(self, t: TypeAliasType) -> Type: