Skip to content

Move some enum and literal related type ops to typeops.py #7909

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 3 additions & 93 deletions mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

import itertools
import fnmatch
import sys
from contextlib import contextmanager

from typing import (
Expand Down Expand Up @@ -74,7 +73,9 @@
from mypy.plugin import Plugin, CheckerPluginInterface
from mypy.sharedparse import BINARY_MAGIC_METHODS
from mypy.scope import Scope
from mypy.typeops import tuple_fallback
from mypy.typeops import (
tuple_fallback, coerce_to_literal, is_singleton_type, try_expanding_enum_to_union
)
from mypy import state, errorcodes as codes
from mypy.traverser import has_return_statement, all_return_statements
from mypy.errorcodes import ErrorCode
Expand Down Expand Up @@ -4766,97 +4767,6 @@ def is_private(node_name: str) -> bool:
return node_name.startswith('__') and not node_name.endswith('__')


def get_enum_values(typ: Instance) -> List[str]:
"""Return the list of values for an Enum."""
return [name for name, sym in typ.type.names.items() if isinstance(sym.node, Var)]


def is_singleton_type(typ: Type) -> bool:
"""Returns 'true' if this type is a "singleton type" -- if there exists
exactly only one runtime value associated with this type.

That is, given two values 'a' and 'b' that have the same type 't',
'is_singleton_type(t)' returns True if and only if the expression 'a is b' is
always true.

Currently, this returns True when given NoneTypes, enum LiteralTypes and
enum types with a single value.

Note that other kinds of LiteralTypes cannot count as singleton types. For
example, suppose we do 'a = 100000 + 1' and 'b = 100001'. It is not guaranteed
that 'a is b' will always be true -- some implementations of Python will end up
constructing two distinct instances of 100001.
"""
typ = get_proper_type(typ)
# TODO: Also make this return True if the type is a bool LiteralType.
# Also make this return True if the type corresponds to ... (ellipsis) or NotImplemented?
return (
isinstance(typ, NoneType) or (isinstance(typ, LiteralType) and typ.is_enum_literal())
or (isinstance(typ, Instance) and typ.type.is_enum and len(get_enum_values(typ)) == 1)
)


def try_expanding_enum_to_union(typ: Type, target_fullname: str) -> ProperType:
"""Attempts to recursively expand any enum Instances with the given target_fullname
into a Union of all of its component LiteralTypes.

For example, if we have:

class Color(Enum):
RED = 1
BLUE = 2
YELLOW = 3

class Status(Enum):
SUCCESS = 1
FAILURE = 2
UNKNOWN = 3

...and if we call `try_expanding_enum_to_union(Union[Color, Status], 'module.Color')`,
this function will return Literal[Color.RED, Color.BLUE, Color.YELLOW, Status].
"""
typ = get_proper_type(typ)

if isinstance(typ, UnionType):
items = [try_expanding_enum_to_union(item, target_fullname) for item in typ.items]
return make_simplified_union(items)
elif isinstance(typ, Instance) and typ.type.is_enum and typ.type.fullname() == target_fullname:
new_items = []
for name, symbol in typ.type.names.items():
if not isinstance(symbol.node, Var):
continue
new_items.append(LiteralType(name, typ))
# SymbolTables are really just dicts, and dicts are guaranteed to preserve
# insertion order only starting with Python 3.7. So, we sort these for older
# versions of Python to help make tests deterministic.
#
# We could probably skip the sort for Python 3.6 since people probably run mypy
# only using CPython, but we might as well for the sake of full correctness.
if sys.version_info < (3, 7):
new_items.sort(key=lambda lit: lit.value)
return make_simplified_union(new_items)
else:
return typ


def coerce_to_literal(typ: Type) -> ProperType:
"""Recursively converts any Instances that have a last_known_value or are
instances of enum types with a single value into the corresponding LiteralType.
"""
typ = get_proper_type(typ)
if isinstance(typ, UnionType):
new_items = [coerce_to_literal(item) for item in typ.items]
return make_simplified_union(new_items)
elif isinstance(typ, Instance):
if typ.last_known_value:
return typ.last_known_value
elif typ.type.is_enum:
enum_values = get_enum_values(typ)
if len(enum_values) == 1:
return LiteralType(value=enum_values[0], fallback=typ)
return typ


def has_bool_item(typ: ProperType) -> bool:
"""Return True if type is 'bool' or a union with a 'bool' item."""
if is_named_instance(typ, 'builtins.bool'):
Expand Down
96 changes: 94 additions & 2 deletions mypy/typeops.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
"""

from typing import cast, Optional, List, Sequence, Set
import sys

from mypy.types import (
TupleType, Instance, FunctionLike, Type, CallableType, TypeVarDef, Overloaded,
Expand All @@ -14,8 +15,8 @@
copy_type
)
from mypy.nodes import (
FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, TypeVar, ARG_STAR, ARG_STAR2, Expression,
StrExpr, ARG_POS
FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, TypeVar, ARG_STAR, ARG_STAR2, ARG_POS,
Expression, StrExpr, Var
)
from mypy.maptype import map_instance_to_supertype
from mypy.expandtype import expand_type_by_instance, expand_type
Expand Down Expand Up @@ -489,3 +490,94 @@ def try_getting_str_literals(expr: Expression, typ: Type) -> Optional[List[str]]
else:
return None
return strings


def get_enum_values(typ: Instance) -> List[str]:
"""Return the list of values for an Enum."""
return [name for name, sym in typ.type.names.items() if isinstance(sym.node, Var)]


def is_singleton_type(typ: Type) -> bool:
"""Returns 'true' if this type is a "singleton type" -- if there exists
exactly only one runtime value associated with this type.

That is, given two values 'a' and 'b' that have the same type 't',
'is_singleton_type(t)' returns True if and only if the expression 'a is b' is
always true.

Currently, this returns True when given NoneTypes, enum LiteralTypes and
enum types with a single value.

Note that other kinds of LiteralTypes cannot count as singleton types. For
example, suppose we do 'a = 100000 + 1' and 'b = 100001'. It is not guaranteed
that 'a is b' will always be true -- some implementations of Python will end up
constructing two distinct instances of 100001.
"""
typ = get_proper_type(typ)
# TODO: Also make this return True if the type is a bool LiteralType.
# Also make this return True if the type corresponds to ... (ellipsis) or NotImplemented?
return (
isinstance(typ, NoneType) or (isinstance(typ, LiteralType) and typ.is_enum_literal())
or (isinstance(typ, Instance) and typ.type.is_enum and len(get_enum_values(typ)) == 1)
)


def try_expanding_enum_to_union(typ: Type, target_fullname: str) -> ProperType:
"""Attempts to recursively expand any enum Instances with the given target_fullname
into a Union of all of its component LiteralTypes.

For example, if we have:

class Color(Enum):
RED = 1
BLUE = 2
YELLOW = 3

class Status(Enum):
SUCCESS = 1
FAILURE = 2
UNKNOWN = 3

...and if we call `try_expanding_enum_to_union(Union[Color, Status], 'module.Color')`,
this function will return Literal[Color.RED, Color.BLUE, Color.YELLOW, Status].
"""
typ = get_proper_type(typ)

if isinstance(typ, UnionType):
items = [try_expanding_enum_to_union(item, target_fullname) for item in typ.items]
return make_simplified_union(items)
elif isinstance(typ, Instance) and typ.type.is_enum and typ.type.fullname() == target_fullname:
new_items = []
for name, symbol in typ.type.names.items():
if not isinstance(symbol.node, Var):
continue
new_items.append(LiteralType(name, typ))
# SymbolTables are really just dicts, and dicts are guaranteed to preserve
# insertion order only starting with Python 3.7. So, we sort these for older
# versions of Python to help make tests deterministic.
#
# We could probably skip the sort for Python 3.6 since people probably run mypy
# only using CPython, but we might as well for the sake of full correctness.
if sys.version_info < (3, 7):
new_items.sort(key=lambda lit: lit.value)
return make_simplified_union(new_items)
else:
return typ


def coerce_to_literal(typ: Type) -> ProperType:
"""Recursively converts any Instances that have a last_known_value or are
instances of enum types with a single value into the corresponding LiteralType.
"""
typ = get_proper_type(typ)
if isinstance(typ, UnionType):
new_items = [coerce_to_literal(item) for item in typ.items]
return make_simplified_union(new_items)
elif isinstance(typ, Instance):
if typ.last_known_value:
return typ.last_known_value
elif typ.type.is_enum:
enum_values = get_enum_values(typ)
if len(enum_values) == 1:
return LiteralType(value=enum_values[0], fallback=typ)
return typ