import sys
from functools import partial
from typing import Callable, List, Optional, Sequence, Tuple, cast
from typing_extensions import Literal
from typing import Type as typing_Type
from mypy import message_registry
from mypy.plugin import Plugin, FunctionContext, MethodContext, CheckerPluginInterface
from mypy.nodes import (
ARG_POS,
ARG_STAR,
TypeInfo,
Context,
FuncDef,
StrExpr,
IntExpr,
Expression,
)
from mypy.types import (
Type,
CallableType,
NoneTyp,
Overloaded,
TypeVarLikeType,
TypeVarType,
Instance,
UnionType,
UninhabitedType,
AnyType,
TypeOfAny,
get_proper_type,
get_proper_types,
)
from mypy.typeops import make_simplified_union
from mypy.checker import TypeChecker
from packaging.version import parse as parse_version
class TrioPlugin(Plugin):
def get_function_hook(
self, fullname: str
) -> Optional[Callable[[FunctionContext], Type]]:
if fullname == "trio_typing.takes_callable_and_args":
return takes_callable_and_args_callback
if fullname == "async_generator.async_generator":
return async_generator_callback
if fullname == "async_generator.yield_":
return yield_callback
if fullname == "async_generator.yield_from_":
return yield_from_callback
return None
class TrioPlugin13(TrioPlugin):
def get_function_hook(
self, fullname: str
) -> Optional[Callable[[FunctionContext], Type]]:
if fullname == "trio_typing.takes_callable_and_args":
return partial(takes_callable_and_args_callback, has_type_var_default=False)
return super().get_function_hook(fullname)
def decode_agen_types_from_return_type(
ctx: FunctionContext, original_async_return_type: Type
) -> Tuple[Type, Type, Type]:
"""Return the yield type, send type, and return type of
an @async_generator decorated function that was
originally declared to return ``original_async_return_type``.
This tries to interpret ``original_async_return_type`` as a union
between the async generator return type (i.e., the thing actually
returned by the decorated function, which becomes the value
associated with a ``StopAsyncIteration`` exception), an optional
``trio_typing.YieldType[X]`` where ``X`` is the type of values
that the async generator yields, and an optional
``trio_typing.SendType[Y]`` where ``Y`` is the type of values that
the async generator expects to be sent. If one of ``YieldType``
and ``SendType`` is specified, the other is assumed to be None;
if neither is specified, both are assumed to be Any.
If ``original_async_return_type`` includes a ``YieldType``
and/or a ``SendType`` but no actual return type, the return
is inferred as ``NoReturn``.
"""
arms = [] # type: Sequence[Type]
resolved_async_return_type = get_proper_type(original_async_return_type)
if isinstance(resolved_async_return_type, UnionType):
arms = resolved_async_return_type.items
else:
arms = [original_async_return_type]
yield_type = None # type: Optional[Type]
send_type = None # type: Optional[Type]
other_arms = [] # type: List[Type]
try:
for orig_arm in arms:
arm = get_proper_type(orig_arm)
if isinstance(arm, Instance):
if arm.type.fullname == "trio_typing.YieldType":
if len(arm.args) != 1:
raise ValueError("YieldType must take one argument")
if yield_type is not None:
raise ValueError("YieldType specified multiple times")
yield_type = arm.args[0]
elif arm.type.fullname == "trio_typing.SendType":
if len(arm.args) != 1:
raise ValueError("SendType must take one argument")
if send_type is not None:
raise ValueError("SendType specified multiple times")
send_type = arm.args[0]
else:
other_arms.append(orig_arm)
else:
other_arms.append(orig_arm)
except ValueError as ex:
ctx.api.fail("invalid @async_generator return type: {}".format(ex), ctx.context)
return (
AnyType(TypeOfAny.from_error),
AnyType(TypeOfAny.from_error),
original_async_return_type,
)
if yield_type is None and send_type is None:
return (
AnyType(TypeOfAny.unannotated),
AnyType(TypeOfAny.unannotated),
original_async_return_type,
)
if yield_type is None:
yield_type = NoneTyp(ctx.context.line, ctx.context.column)
if send_type is None:
send_type = NoneTyp(ctx.context.line, ctx.context.column)
if not other_arms:
return (
yield_type,
send_type,
UninhabitedType(
is_noreturn=True, line=ctx.context.line, column=ctx.context.column
),
)
else:
return (
yield_type,
send_type,
make_simplified_union(other_arms, ctx.context.line, ctx.context.column),
)
def async_generator_callback(ctx: FunctionContext) -> Type:
"""Handle @async_generator.
This moves the yield type and send type declarations from
the return type of the decorated function to the appropriate
type parameters of ``trio_typing.CompatAsyncGenerator``.
That is, if you say::
@async_generator
async def example() -> Union[str, YieldType[bool], SendType[int]]:
...
then the decorated ``example()`` will return values of type
``CompatAsyncGenerator[bool, int, str]``, as opposed to
``CompatAsyncGenerator[Any, Any, Union[str,
YieldType[bool], SendType[int]]`` without the plugin.
"""
decorator_return_type = ctx.default_return_type
if not isinstance(decorator_return_type, CallableType):
return decorator_return_type
agen_return_type = get_proper_type(decorator_return_type.ret_type)
if (
isinstance(agen_return_type, Instance)
and agen_return_type.type.fullname == "trio_typing.CompatAsyncGenerator"
and len(agen_return_type.args) == 3
):
return decorator_return_type.copy_modified(
ret_type=agen_return_type.copy_modified(
args=list(
decode_agen_types_from_return_type(ctx, agen_return_type.args[2])
)
)
)
return decorator_return_type
def decode_enclosing_agen_types(ctx: FunctionContext) -> Tuple[Type, Type]:
"""Return the yield and send types that would be returned by
decode_agen_types_from_return_type() for the function that's
currently being typechecked, i.e., the function that contains the
call described in ``ctx``.
"""
private_api = cast(TypeChecker, ctx.api)
enclosing_func = private_api.scope.top_function()
if (
enclosing_func is None
or not isinstance(enclosing_func, FuncDef)
or not enclosing_func.is_coroutine
or not enclosing_func.is_decorated
):
# we can't actually detect the @async_generator decorator but
# we'll at least notice if it couldn't possibly be present
ctx.api.fail(
"async_generator.yield_() outside an @async_generator func", ctx.context
)
return AnyType(TypeOfAny.from_error), AnyType(TypeOfAny.from_error)
# The enclosing function type Callable[...] and its return type
# Coroutine[...] were both produced by mypy, rather than typed by
# the user, so they can't be type aliases; thus there's no need to
# use get_proper_type() here.
if (
isinstance(enclosing_func.type, CallableType)
and isinstance(enclosing_func.type.ret_type, Instance)
and enclosing_func.type.ret_type.type.fullname == "typing.Coroutine"
and len(enclosing_func.type.ret_type.args) == 3
):
yield_type, send_type, _ = decode_agen_types_from_return_type(
ctx, enclosing_func.type.ret_type.args[2]
)
return yield_type, send_type
return (
AnyType(TypeOfAny.implementation_artifact),
AnyType(TypeOfAny.implementation_artifact),
)
def yield_callback(ctx: FunctionContext) -> Type:
"""Provide a more specific argument and return type for yield_()
inside an @async_generator.
"""
if len(ctx.arg_types) == 0:
arg_type = NoneTyp(ctx.context.line, ctx.context.column) # type: Type
elif ctx.arg_types and len(ctx.arg_types[0]) == 1:
arg_type = ctx.arg_types[0][0]
else:
return ctx.default_return_type
private_api = cast(TypeChecker, ctx.api)
yield_type, send_type = decode_enclosing_agen_types(ctx)
if yield_type is not None and send_type is not None:
private_api.check_subtype(
subtype=arg_type,
supertype=yield_type,
context=ctx.context,
msg=message_registry.INCOMPATIBLE_TYPES,
subtype_label="yield_ argument",
supertype_label="declared YieldType",
)
return ctx.api.named_generic_type("typing.Awaitable", [send_type])
return ctx.default_return_type
def yield_from_callback(ctx: FunctionContext) -> Type:
"""Provide a better typecheck for yield_from_()."""
if ctx.arg_types and len(ctx.arg_types[0]) == 1:
arg_type = get_proper_type(ctx.arg_types[0][0])
else:
return ctx.default_return_type
private_api = cast(TypeChecker, ctx.api)
our_yield_type, our_send_type = decode_enclosing_agen_types(ctx)
if our_yield_type is None or our_send_type is None:
return ctx.default_return_type
if (
isinstance(arg_type, Instance)
and arg_type.type.fullname
in (
"trio_typing.CompatAsyncGenerator",
"trio_typing.AsyncGenerator",
"typing.AsyncGenerator",
)
and len(arg_type.args) >= 2
):
their_yield_type, their_send_type = arg_type.args[:2]
private_api.check_subtype(
subtype=their_yield_type,
supertype=our_yield_type,
context=ctx.context,
msg=message_registry.INCOMPATIBLE_TYPES,
subtype_label="yield_from_ argument YieldType",
supertype_label="local declared YieldType",
)
private_api.check_subtype(
subtype=our_send_type,
supertype=their_send_type,
context=ctx.context,
msg=message_registry.INCOMPATIBLE_TYPES,
subtype_label="local declared SendType",
supertype_label="yield_from_ argument SendType",
)
elif isinstance(arg_type, Instance):
private_api.check_subtype(
subtype=arg_type,
supertype=ctx.api.named_generic_type(
"typing.AsyncIterable", [our_yield_type]
),
context=ctx.context,
msg=message_registry.INCOMPATIBLE_TYPES,
subtype_label="yield_from_ argument type",
supertype_label="expected iterable type",
)
return ctx.default_return_type
def takes_callable_and_args_callback(
ctx: FunctionContext, has_type_var_default: bool = True
) -> Type:
"""Automate the boilerplate for writing functions that accept
arbitrary positional arguments of the same type accepted by
a callable.
For example, this supports writing::
@trio_typing.takes_callable_and_args
def start_soon(
self,
async_fn: Callable[[VarArg()], Any],
*args: Any,
) -> None: ...
instead of::
T1 = TypeVar("T1")
T2 = TypeVar("T2")
T3 = TypeVar("T3")
T4 = TypeVar("T4")
@overload
def start_soon(
self,
async_fn: Callable[[], Any],
) -> None: ...
@overload
def start_soon(
self,
async_fn: Callable[[T1], Any],
__arg1: T1,
) -> None: ...
@overload
def start_soon(
self,
async_fn: Callable[[T1, T2], Any],
__arg1: T1,
__arg2: T2
) -> None: ...
# etc
"""
try:
if not ctx.arg_types or len(ctx.arg_types[0]) != 1:
raise ValueError("must be used as a decorator")
fn_type = get_proper_type(ctx.arg_types[0][0])
if not isinstance(fn_type, CallableType) or not isinstance(
get_proper_type(ctx.default_return_type), CallableType
):
raise ValueError("must be used as a decorator")
callable_idx = -1 # index in function arguments of the callable
callable_args_idx = -1 # index in callable arguments of the StarArgs
callable_ty = None # type: Optional[CallableType]
args_idx = -1 # index in function arguments of the StarArgs
for idx, (kind, ty) in enumerate(zip(fn_type.arg_kinds, fn_type.arg_types)):
ty = get_proper_type(ty)
if isinstance(ty, AnyType) and kind == ARG_STAR:
assert args_idx == -1
args_idx = idx
elif isinstance(ty, (UnionType, CallableType)) and kind == ARG_POS:
# turn Union[Callable[..., T], Callable[[VarArg()], T]]
# into Callable[[VarArg()], T]
# (the union makes it not fail when the plugin is not being used)
if isinstance(ty, UnionType):
for arm in get_proper_types(ty.items):
if (
isinstance(arm, CallableType)
and not arm.is_ellipsis_args
and any(kind_ == ARG_STAR for kind_ in arm.arg_kinds)
):
ty = arm
break
else:
continue
for idx_, (kind_, ty_) in enumerate(zip(ty.arg_kinds, ty.arg_types)):
ty_ = get_proper_type(ty_)
if isinstance(ty_, AnyType) and kind_ == ARG_STAR:
if callable_idx != -1:
raise ValueError(
"decorated function may only take one callable "
"that has an argument of type VarArg()"
)
callable_idx = idx
callable_args_idx = idx_
callable_ty = ty
if args_idx == -1:
raise ValueError("decorated function must take *args: Any")
if callable_idx == -1 or callable_ty is None:
raise ValueError(
"decorated function must take a callable that has a "
"argument of type mypy_extensions.VarArg()"
)
expanded_fns = [] # type: List[CallableType]
type_var_types = [] # type: List[TypeVarType]
for arg_idx in range(1, 7): # provides overloads for 0 through 5 arguments
arg_types = list(fn_type.arg_types)
arg_types[callable_idx] = callable_ty.copy_modified(
arg_types=(
callable_ty.arg_types[:callable_args_idx]
+ cast(List[Type], type_var_types)
+ callable_ty.arg_types[callable_args_idx + 1 :]
),
arg_kinds=(
callable_ty.arg_kinds[:callable_args_idx]
+ ([ARG_POS] * len(type_var_types))
+ callable_ty.arg_kinds[callable_args_idx + 1 :]
),
arg_names=(
callable_ty.arg_names[:callable_args_idx]
+ ([None] * len(type_var_types))
+ callable_ty.arg_names[callable_args_idx + 1 :]
),
# Note that we do *not* append `type_var_types` to
# `callable_ty.variables`. Even though `*type_var_types` are in our new
# `callable_ty`'s argument types, they are *not* type variables that get
# bound when our new `callable_ty` gets called. They get bound when the
# `expanded_fn` that references our new `callable_ty` gets called.
)
expanded_fns.append(
fn_type.copy_modified(
arg_types=(
arg_types[:args_idx]
+ cast(List[Type], type_var_types)
+ arg_types[args_idx + 1 :]
),
arg_kinds=(
fn_type.arg_kinds[:args_idx]
+ ([ARG_POS] * len(type_var_types))
+ fn_type.arg_kinds[args_idx + 1 :]
),
arg_names=(
fn_type.arg_names[:args_idx]
+ ([None] * len(type_var_types))
+ fn_type.arg_names[args_idx + 1 :]
),
variables=(
list(fn_type.variables)
+ cast(List[TypeVarLikeType], type_var_types)
),
)
)
if has_type_var_default:
type_var_types.append(
TypeVarType(
"__T{}".format(arg_idx),
"__T{}".format(arg_idx),
-len(fn_type.variables) - arg_idx - 1,
[],
ctx.api.named_generic_type("builtins.object", []),
line=ctx.context.line,
column=ctx.context.column,
default=AnyType(TypeOfAny.from_omitted_generics),
)
)
else:
type_var_types.append(
TypeVarType(
"__T{}".format(arg_idx),
"__T{}".format(arg_idx),
-len(fn_type.variables) - arg_idx - 1,
[],
ctx.api.named_generic_type("builtins.object", []),
line=ctx.context.line,
column=ctx.context.column,
) # type: ignore[call-arg]
)
return Overloaded(expanded_fns)
except ValueError as ex:
ctx.api.fail(
"invalid use of @takes_callable_and_args: {}".format(ex), ctx.context
)
return ctx.default_return_type
def plugin(version: str) -> typing_Type[Plugin]:
mypy_version = parse_version(version)
if mypy_version < parse_version("1.0"):
raise RuntimeError("This version of trio-typing requires at least mypy 1.0.")
elif mypy_version < parse_version("1.4"):
return TrioPlugin13
else:
return TrioPlugin