mypy-dev 1.9.0.dev0__py3-none-any.whl

mypy/plugins/enums.py

@@ -0,0 +1,259 @@
+"""
+This file contains a variety of plugins for refining how mypy infers types of
+expressions involving Enums.
+
+Currently, this file focuses on providing better inference for expressions like
+'SomeEnum.FOO.name' and 'SomeEnum.FOO.value'. Note that the type of both expressions
+will vary depending on exactly which instance of SomeEnum we're looking at.
+
+Note that this file does *not* contain all special-cased logic related to enums:
+we actually bake some of it directly in to the semantic analysis layer (see
+semanal_enum.py).
+"""
+
+from __future__ import annotations
+
+from typing import Final, Iterable, Sequence, TypeVar, cast
+
+import mypy.plugin  # To avoid circular imports.
+from mypy.nodes import TypeInfo
+from mypy.semanal_enum import ENUM_BASES
+from mypy.subtypes import is_equivalent
+from mypy.typeops import fixup_partial_type, make_simplified_union
+from mypy.types import CallableType, Instance, LiteralType, ProperType, Type, get_proper_type
+
+ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | {
+    f"{prefix}._name_" for prefix in ENUM_BASES
+}
+ENUM_VALUE_ACCESS: Final = {f"{prefix}.value" for prefix in ENUM_BASES} | {
+    f"{prefix}._value_" for prefix in ENUM_BASES
+}
+
+
+def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> Type:
+    """This plugin refines the 'name' attribute in enums to act as if
+    they were declared to be final.
+
+    For example, the expression 'MyEnum.FOO.name' normally is inferred
+    to be of type 'str'.
+
+    This plugin will instead make the inferred type be a 'str' where the
+    last known value is 'Literal["FOO"]'. This means it would be legal to
+    use 'MyEnum.FOO.name' in contexts that expect a Literal type, just like
+    any other Final variable or attribute.
+
+    This plugin assumes that the provided context is an attribute access
+    matching one of the strings found in 'ENUM_NAME_ACCESS'.
+    """
+    enum_field_name = _extract_underlying_field_name(ctx.type)
+    if enum_field_name is None:
+        return ctx.default_attr_type
+    else:
+        str_type = ctx.api.named_generic_type("builtins.str", [])
+        literal_type = LiteralType(enum_field_name, fallback=str_type)
+        return str_type.copy_modified(last_known_value=literal_type)
+
+
+_T = TypeVar("_T")
+
+
+def _first(it: Iterable[_T]) -> _T | None:
+    """Return the first value from any iterable.
+
+    Returns ``None`` if the iterable is empty.
+    """
+    for val in it:
+        return val
+    return None
+
+
+def _infer_value_type_with_auto_fallback(
+    ctx: mypy.plugin.AttributeContext, proper_type: ProperType | None
+) -> Type | None:
+    """Figure out the type of an enum value accounting for `auto()`.
+
+    This method is a no-op for a `None` proper_type and also in the case where
+    the type is not "enum.auto"
+    """
+    if proper_type is None:
+        return None
+    proper_type = get_proper_type(fixup_partial_type(proper_type))
+    if not (isinstance(proper_type, Instance) and proper_type.type.fullname == "enum.auto"):
+        return proper_type
+    assert isinstance(ctx.type, Instance), "An incorrect ctx.type was passed."
+    info = ctx.type.type
+    # Find the first _generate_next_value_ on the mro.  We need to know
+    # if it is `Enum` because `Enum` types say that the return-value of
+    # `_generate_next_value_` is `Any`.  In reality the default `auto()`
+    # returns an `int` (presumably the `Any` in typeshed is to make it
+    # easier to subclass and change the returned type).
+    type_with_gnv = _first(ti for ti in info.mro if ti.names.get("_generate_next_value_"))
+    if type_with_gnv is None:
+        return ctx.default_attr_type
+
+    stnode = type_with_gnv.names["_generate_next_value_"]
+
+    # This should be a `CallableType`
+    node_type = get_proper_type(stnode.type)
+    if isinstance(node_type, CallableType):
+        if type_with_gnv.fullname == "enum.Enum":
+            int_type = ctx.api.named_generic_type("builtins.int", [])
+            return int_type
+        return get_proper_type(node_type.ret_type)
+    return ctx.default_attr_type
+
+
+def _implements_new(info: TypeInfo) -> bool:
+    """Check whether __new__ comes from enum.Enum or was implemented in a
+    subclass. In the latter case, we must infer Any as long as mypy can't infer
+    the type of _value_ from assignments in __new__.
+    """
+    type_with_new = _first(
+        ti
+        for ti in info.mro
+        if ti.names.get("__new__") and not ti.fullname.startswith("builtins.")
+    )
+    if type_with_new is None:
+        return False
+    return type_with_new.fullname not in ("enum.Enum", "enum.IntEnum", "enum.StrEnum")
+
+
+def enum_value_callback(ctx: mypy.plugin.AttributeContext) -> Type:
+    """This plugin refines the 'value' attribute in enums to refer to
+    the original underlying value. For example, suppose we have the
+    following:
+
+        class SomeEnum:
+            FOO = A()
+            BAR = B()
+
+    By default, mypy will infer that 'SomeEnum.FOO.value' and
+    'SomeEnum.BAR.value' both are of type 'Any'. This plugin refines
+    this inference so that mypy understands the expressions are
+    actually of types 'A' and 'B' respectively. This better reflects
+    the actual runtime behavior.
+
+    This plugin works simply by looking up the original value assigned
+    to the enum. For example, when this plugin sees 'SomeEnum.BAR.value',
+    it will look up whatever type 'BAR' had in the SomeEnum TypeInfo and
+    use that as the inferred type of the overall expression.
+
+    This plugin assumes that the provided context is an attribute access
+    matching one of the strings found in 'ENUM_VALUE_ACCESS'.
+    """
+    enum_field_name = _extract_underlying_field_name(ctx.type)
+    if enum_field_name is None:
+        # We do not know the enum field name (perhaps it was passed to a
+        # function and we only know that it _is_ a member).  All is not lost
+        # however, if we can prove that the all of the enum members have the
+        # same value-type, then it doesn't matter which member was passed in.
+        # The value-type is still known.
+        if isinstance(ctx.type, Instance):
+            info = ctx.type.type
+
+            # As long as mypy doesn't understand attribute creation in __new__,
+            # there is no way to predict the value type if the enum class has a
+            # custom implementation
+            if _implements_new(info):
+                return ctx.default_attr_type
+
+            stnodes = (info.get(name) for name in info.names)
+
+            # Enums _can_ have methods and instance attributes.
+            # Omit methods and attributes created by assigning to self.*
+            # for our value inference.
+            node_types = (
+                get_proper_type(n.type) if n else None
+                for n in stnodes
+                if n is None or not n.implicit
+            )
+            proper_types = [
+                _infer_value_type_with_auto_fallback(ctx, t)
+                for t in node_types
+                if t is None or not isinstance(t, CallableType)
+            ]
+            underlying_type = _first(proper_types)
+            if underlying_type is None:
+                return ctx.default_attr_type
+
+            # At first we try to predict future `value` type if all other items
+            # have the same type. For example, `int`.
+            # If this is the case, we simply return this type.
+            # See https://github.com/python/mypy/pull/9443
+            all_same_value_type = all(
+                proper_type is not None and proper_type == underlying_type
+                for proper_type in proper_types
+            )
+            if all_same_value_type:
+                if underlying_type is not None:
+                    return underlying_type
+
+            # But, after we started treating all `Enum` values as `Final`,
+            # we start to infer types in
+            # `item = 1` as `Literal[1]`, not just `int`.
+            # So, for example types in this `Enum` will all be different:
+            #
+            #  class Ordering(IntEnum):
+            #      one = 1
+            #      two = 2
+            #      three = 3
+            #
+            # We will infer three `Literal` types here.
+            # They are not the same, but they are equivalent.
+            # So, we unify them to make sure `.value` prediction still works.
+            # Result will be `Literal[1] | Literal[2] | Literal[3]` for this case.
+            all_equivalent_types = all(
+                proper_type is not None and is_equivalent(proper_type, underlying_type)
+                for proper_type in proper_types
+            )
+            if all_equivalent_types:
+                return make_simplified_union(cast(Sequence[Type], proper_types))
+        return ctx.default_attr_type
+
+    assert isinstance(ctx.type, Instance)
+    info = ctx.type.type
+
+    # As long as mypy doesn't understand attribute creation in __new__,
+    # there is no way to predict the value type if the enum class has a
+    # custom implementation
+    if _implements_new(info):
+        return ctx.default_attr_type
+
+    stnode = info.get(enum_field_name)
+    if stnode is None:
+        return ctx.default_attr_type
+
+    underlying_type = _infer_value_type_with_auto_fallback(ctx, get_proper_type(stnode.type))
+    if underlying_type is None:
+        return ctx.default_attr_type
+
+    return underlying_type
+
+
+def _extract_underlying_field_name(typ: Type) -> str | None:
+    """If the given type corresponds to some Enum instance, returns the
+    original name of that enum. For example, if we receive in the type
+    corresponding to 'SomeEnum.FOO', we return the string "SomeEnum.Foo".
+
+    This helper takes advantage of the fact that Enum instances are valid
+    to use inside Literal[...] types. An expression like 'SomeEnum.FOO' is
+    actually represented by an Instance type with a Literal enum fallback.
+
+    We can examine this Literal fallback to retrieve the string.
+    """
+    typ = get_proper_type(typ)
+    if not isinstance(typ, Instance):
+        return None
+
+    if not typ.type.is_enum:
+        return None
+
+    underlying_literal = typ.last_known_value
+    if underlying_literal is None:
+        return None
+
+    # The checks above have verified this LiteralType is representing an enum value,
+    # which means the 'value' field is guaranteed to be the name of the enum field
+    # as a string.
+    assert isinstance(underlying_literal.value, str)
+    return underlying_literal.value

mypy/plugins/functools.py

@@ -0,0 +1,104 @@
+"""Plugin for supporting the functools standard library module."""
+
+from __future__ import annotations
+
+from typing import Final, NamedTuple
+
+import mypy.plugin
+from mypy.nodes import ARG_POS, ARG_STAR2, Argument, FuncItem, Var
+from mypy.plugins.common import add_method_to_class
+from mypy.types import AnyType, CallableType, Type, TypeOfAny, UnboundType, get_proper_type
+
+functools_total_ordering_makers: Final = {"functools.total_ordering"}
+
+_ORDERING_METHODS: Final = {"__lt__", "__le__", "__gt__", "__ge__"}
+
+
+class _MethodInfo(NamedTuple):
+    is_static: bool
+    type: CallableType
+
+
+def functools_total_ordering_maker_callback(
+    ctx: mypy.plugin.ClassDefContext, auto_attribs_default: bool = False
+) -> bool:
+    """Add dunder methods to classes decorated with functools.total_ordering."""
+    comparison_methods = _analyze_class(ctx)
+    if not comparison_methods:
+        ctx.api.fail(
+            'No ordering operation defined when using "functools.total_ordering": < > <= >=',
+            ctx.reason,
+        )
+        return True
+
+    # prefer __lt__ to __le__ to __gt__ to __ge__
+    root = max(comparison_methods, key=lambda k: (comparison_methods[k] is None, k))
+    root_method = comparison_methods[root]
+    if not root_method:
+        # None of the defined comparison methods can be analysed
+        return True
+
+    other_type = _find_other_type(root_method)
+    bool_type = ctx.api.named_type("builtins.bool")
+    ret_type: Type = bool_type
+    if root_method.type.ret_type != ctx.api.named_type("builtins.bool"):
+        proper_ret_type = get_proper_type(root_method.type.ret_type)
+        if not (
+            isinstance(proper_ret_type, UnboundType)
+            and proper_ret_type.name.split(".")[-1] == "bool"
+        ):
+            ret_type = AnyType(TypeOfAny.implementation_artifact)
+    for additional_op in _ORDERING_METHODS:
+        # Either the method is not implemented
+        # or has an unknown signature that we can now extrapolate.
+        if not comparison_methods.get(additional_op):
+            args = [Argument(Var("other", other_type), other_type, None, ARG_POS)]
+            add_method_to_class(ctx.api, ctx.cls, additional_op, args, ret_type)
+
+    return True
+
+
+def _find_other_type(method: _MethodInfo) -> Type:
+    """Find the type of the ``other`` argument in a comparison method."""
+    first_arg_pos = 0 if method.is_static else 1
+    cur_pos_arg = 0
+    other_arg = None
+    for arg_kind, arg_type in zip(method.type.arg_kinds, method.type.arg_types):
+        if arg_kind.is_positional():
+            if cur_pos_arg == first_arg_pos:
+                other_arg = arg_type
+                break
+
+            cur_pos_arg += 1
+        elif arg_kind != ARG_STAR2:
+            other_arg = arg_type
+            break
+
+    if other_arg is None:
+        return AnyType(TypeOfAny.implementation_artifact)
+
+    return other_arg
+
+
+def _analyze_class(ctx: mypy.plugin.ClassDefContext) -> dict[str, _MethodInfo | None]:
+    """Analyze the class body, its parents, and return the comparison methods found."""
+    # Traverse the MRO and collect ordering methods.
+    comparison_methods: dict[str, _MethodInfo | None] = {}
+    # Skip object because total_ordering does not use methods from object
+    for cls in ctx.cls.info.mro[:-1]:
+        for name in _ORDERING_METHODS:
+            if name in cls.names and name not in comparison_methods:
+                node = cls.names[name].node
+                if isinstance(node, FuncItem) and isinstance(node.type, CallableType):
+                    comparison_methods[name] = _MethodInfo(node.is_static, node.type)
+                    continue
+
+                if isinstance(node, Var):
+                    proper_type = get_proper_type(node.type)
+                    if isinstance(proper_type, CallableType):
+                        comparison_methods[name] = _MethodInfo(node.is_staticmethod, proper_type)
+                        continue
+
+                comparison_methods[name] = None
+
+    return comparison_methods

mypy/plugins/proper_plugin.py

@@ -0,0 +1,175 @@
+"""
+This plugin is helpful for mypy development itself.
+By default, it is not enabled for mypy users.
+
+It also can be used by plugin developers as a part of their CI checks.
+
+It finds missing ``get_proper_type()`` call, which can lead to multiple errors.
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo
+from mypy.plugin import FunctionContext, Plugin
+from mypy.subtypes import is_proper_subtype
+from mypy.types import (
+    AnyType,
+    FunctionLike,
+    Instance,
+    NoneTyp,
+    ProperType,
+    TupleType,
+    Type,
+    UnionType,
+    get_proper_type,
+    get_proper_types,
+)
+
+
+class ProperTypePlugin(Plugin):
+    """
+    A plugin to ensure that every type is expanded before doing any special-casing.
+
+    This solves the problem that we have hundreds of call sites like:
+
+        if isinstance(typ, UnionType):
+            ...  # special-case union
+
+    But after introducing a new type TypeAliasType (and removing immediate expansion)
+    all these became dangerous because typ may be e.g. an alias to union.
+    """
+
+    def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
+        if fullname == "builtins.isinstance":
+            return isinstance_proper_hook
+        if fullname == "mypy.types.get_proper_type":
+            return proper_type_hook
+        if fullname == "mypy.types.get_proper_types":
+            return proper_types_hook
+        return None
+
+
+def isinstance_proper_hook(ctx: FunctionContext) -> Type:
+    if len(ctx.arg_types) != 2 or not ctx.arg_types[1]:
+        return ctx.default_return_type
+
+    right = get_proper_type(ctx.arg_types[1][0])
+    for arg in ctx.arg_types[0]:
+        if (
+            is_improper_type(arg) or isinstance(get_proper_type(arg), AnyType)
+        ) and is_dangerous_target(right):
+            if is_special_target(right):
+                return ctx.default_return_type
+            ctx.api.fail(
+                "Never apply isinstance() to unexpanded types;"
+                " use mypy.types.get_proper_type() first",
+                ctx.context,
+            )
+            ctx.api.note(  # type: ignore[attr-defined]
+                "If you pass on the original type"
+                " after the check, always use its unexpanded version",
+                ctx.context,
+            )
+    return ctx.default_return_type
+
+
+def is_special_target(right: ProperType) -> bool:
+    """Whitelist some special cases for use in isinstance() with improper types."""
+    if isinstance(right, FunctionLike) and right.is_type_obj():
+        if right.type_object().fullname == "builtins.tuple":
+            # Used with Union[Type, Tuple[Type, ...]].
+            return True
+        if right.type_object().fullname in (
+            "mypy.types.Type",
+            "mypy.types.ProperType",
+            "mypy.types.TypeAliasType",
+        ):
+            # Special case: things like assert isinstance(typ, ProperType) are always OK.
+            return True
+        if right.type_object().fullname in (
+            "mypy.types.UnboundType",
+            "mypy.types.TypeVarLikeType",
+            "mypy.types.TypeVarType",
+            "mypy.types.UnpackType",
+            "mypy.types.TypeVarTupleType",
+            "mypy.types.ParamSpecType",
+            "mypy.types.Parameters",
+            "mypy.types.RawExpressionType",
+            "mypy.types.EllipsisType",
+            "mypy.types.StarType",
+            "mypy.types.TypeList",
+            "mypy.types.CallableArgument",
+            "mypy.types.PartialType",
+            "mypy.types.ErasedType",
+            "mypy.types.DeletedType",
+            "mypy.types.RequiredType",
+        ):
+            # Special case: these are not valid targets for a type alias and thus safe.
+            # TODO: introduce a SyntheticType base to simplify this?
+            return True
+    elif isinstance(right, TupleType):
+        return all(is_special_target(t) for t in get_proper_types(right.items))
+    return False
+
+
+def is_improper_type(typ: Type) -> bool:
+    """Is this a type that is not a subtype of ProperType?"""
+    typ = get_proper_type(typ)
+    if isinstance(typ, Instance):
+        info = typ.type
+        return info.has_base("mypy.types.Type") and not info.has_base("mypy.types.ProperType")
+    if isinstance(typ, UnionType):
+        return any(is_improper_type(t) for t in typ.items)
+    return False
+
+
+def is_dangerous_target(typ: ProperType) -> bool:
+    """Is this a dangerous target (right argument) for an isinstance() check?"""
+    if isinstance(typ, TupleType):
+        return any(is_dangerous_target(get_proper_type(t)) for t in typ.items)
+    if isinstance(typ, FunctionLike) and typ.is_type_obj():
+        return typ.type_object().has_base("mypy.types.Type")
+    return False
+
+
+def proper_type_hook(ctx: FunctionContext) -> Type:
+    """Check if this get_proper_type() call is not redundant."""
+    arg_types = ctx.arg_types[0]
+    if arg_types:
+        arg_type = get_proper_type(arg_types[0])
+        proper_type = get_proper_type_instance(ctx)
+        if is_proper_subtype(arg_type, UnionType.make_union([NoneTyp(), proper_type])):
+            # Minimize amount of spurious errors from overload machinery.
+            # TODO: call the hook on the overload as a whole?
+            if isinstance(arg_type, (UnionType, Instance)):
+                ctx.api.fail("Redundant call to get_proper_type()", ctx.context)
+    return ctx.default_return_type
+
+
+def proper_types_hook(ctx: FunctionContext) -> Type:
+    """Check if this get_proper_types() call is not redundant."""
+    arg_types = ctx.arg_types[0]
+    if arg_types:
+        arg_type = arg_types[0]
+        proper_type = get_proper_type_instance(ctx)
+        item_type = UnionType.make_union([NoneTyp(), proper_type])
+        ok_type = ctx.api.named_generic_type("typing.Iterable", [item_type])
+        if is_proper_subtype(arg_type, ok_type):
+            ctx.api.fail("Redundant call to get_proper_types()", ctx.context)
+    return ctx.default_return_type
+
+
+def get_proper_type_instance(ctx: FunctionContext) -> Instance:
+    checker = ctx.api
+    assert isinstance(checker, TypeChecker)
+    types = checker.modules["mypy.types"]
+    proper_type_info = types.names["ProperType"]
+    assert isinstance(proper_type_info.node, TypeInfo)
+    return Instance(proper_type_info.node, [])
+
+
+def plugin(version: str) -> type[ProperTypePlugin]:
+    return ProperTypePlugin