From b6b2f19c04c6ece4b9f25d910f18d7ca489716d3 Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Fri, 8 Nov 2019 12:04:41 -0800 Subject: [PATCH] Move some enum and literal related type ops to typeops.py 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. --- mypy/checker.py | 96 ++----------------------------------------------- mypy/typeops.py | 96 +++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 97 insertions(+), 95 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 320a2cf83f2b..f426b2ecf582 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -2,7 +2,6 @@ import itertools import fnmatch -import sys from contextlib import contextmanager from typing import ( @@ -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 @@ -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'): diff --git a/mypy/typeops.py b/mypy/typeops.py index 8db2158d809c..acac6cedcf1b 100644 --- a/mypy/typeops.py +++ b/mypy/typeops.py @@ -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, @@ -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 @@ -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