Source code for returns.contrib.mypy._features.do_notation

from typing import Final, Optional, Tuple

from mypy.maptype import map_instance_to_supertype
from mypy.nodes import Expression, GeneratorExpr, TypeInfo
from mypy.plugin import MethodContext
from mypy.subtypes import is_subtype
from mypy.typeops import make_simplified_union
from mypy.types import AnyType, CallableType, Instance
from mypy.types import Type as MypyType
from mypy.types import TypeOfAny, UnionType, get_proper_type

    'Invalid type supplied in do-notation: expected "{0}", got "{1}"'
    'Literal generator expression is required, not a variable or function call'
    'Using "if" conditions inside a generator is not allowed'

[docs]def analyze(ctx: MethodContext) -> MypyType: """ Used to handle validation and error types in :ref:`do-notation`. What it does? 1. For all types we ensure that only a single container type is used in a single do-notation. We don't allow mixing them. 2. For types with error types (like ``Result``), it inferes what possible errors types can we have. The result is a ``Union`` of all possible errors. 3. Ensures that expression passed into ``.do`` method is literal. 4. Checks that default value is provided if generator expression has ``if`` conditions inside. """ if not ctx.args or not ctx.args[0]: return ctx.default_return_type expr = ctx.args[0][0] if not isinstance(expr, GeneratorExpr):, expr) return ctx.default_return_type if not isinstance(ctx.type, CallableType): return ctx.default_return_type if not isinstance(ctx.default_return_type, Instance): return ctx.default_return_type return _do_notation( expr=expr, type_info=ctx.type.type_object(), default_return_type=ctx.default_return_type, ctx=ctx, )
def _do_notation( expr: GeneratorExpr, type_info: TypeInfo, default_return_type: Instance, ctx: MethodContext, ) -> MypyType: types = [] for seq in expr.sequences: error_type = _try_fetch_error_type(type_info, seq, ctx) if error_type is not None: types.append(error_type) _check_if_conditions(expr, ctx) if types: return default_return_type.copy_modified( args=[ default_return_type.args[0], make_simplified_union(types), *default_return_type.args[2:], ], ) return default_return_type def _try_fetch_error_type( type_info: TypeInfo, seq: Expression, ctx: MethodContext, ) -> Optional[MypyType]: inst = Instance( type_info, [ AnyType(TypeOfAny.implementation_artifact) for _ in type_info.type_vars ], ) typ = ctx.api.expr_checker.accept(seq) # type: ignore if is_subtype(typ, inst, ignore_type_params=True): is_success, error_type = _extract_error_type(typ, type_info) if is_success: return error_type _INVALID_DO_NOTATION_SOURCE.format(inst, typ), seq, ) return None def _extract_error_type( typ: MypyType, type_info: TypeInfo, ) -> Tuple[bool, Optional[MypyType]]: typ = get_proper_type(typ) if isinstance(typ, Instance): return True, _decide_error_type( map_instance_to_supertype(typ, type_info), ) if isinstance(typ, UnionType): types = [] is_success = True for type_item in typ.items: is_success, error_type = _extract_error_type(type_item, type_info) if error_type is not None: types.append(error_type) if is_success: return True, make_simplified_union(types) return False, None def _decide_error_type(typ: Instance) -> Optional[MypyType]: if len(typ.args) < 2: return None if isinstance(get_proper_type(typ.args[1]), AnyType): return None return typ.args[1] def _check_if_conditions( expr: GeneratorExpr, ctx: MethodContext, ) -> None: if any(cond for cond in expr.condlists):, expr)