Skip to content

Commit

Permalink
Move some enum and literal related type ops to typeops.py (python#7909)
Browse files Browse the repository at this point in the history
I keep wanting access to these functions outside of checker.py
in various diffs I'm working on.

So as suggested earlier, I might as well just do this refactoring
now, independent of the other PRs.
  • Loading branch information
Michael0x2a authored Nov 8, 2019
1 parent 59617e8 commit 84126ab
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 95 deletions.
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 @@ -73,7 +72,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 @@ -4771,97 +4772,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, TypeAliasType
)
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 @@ -495,3 +496,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

0 comments on commit 84126ab

Please sign in to comment.