Skip to content

Commit 86accbb

Browse files
author
Guido van Rossum
committed
Support functional API for Enum.
Fixes #2306.
1 parent ed7d0c0 commit 86accbb

8 files changed

Lines changed: 153 additions & 5 deletions

File tree

‎mypy/checker.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
YieldFromExpr, NamedTupleExpr, TypedDictExpr, SetComprehension,
2525
DictionaryComprehension, ComplexExpr, EllipsisExpr, TypeAliasExpr,
2626
RefExpr, YieldExpr, BackquoteExpr, ImportFrom, ImportAll, ImportBase,
27-
AwaitExpr, PromoteExpr, Node,
27+
AwaitExpr, PromoteExpr, Node, EnumCallExpr,
2828
ARG_POS, MDEF,
2929
CONTRAVARIANT, COVARIANT)
3030
from mypy import nodes
@@ -2244,6 +2244,9 @@ def visit_newtype_expr(self, e: NewTypeExpr) -> Type:
22442244
def visit_namedtuple_expr(self, e: NamedTupleExpr) -> Type:
22452245
return self.expr_checker.visit_namedtuple_expr(e)
22462246

2247+
def visit_enum_call_expr(self, e: EnumCallExpr) -> Type:
2248+
return self.expr_checker.visit_enum_call_expr(e)
2249+
22472250
def visit_typeddict_expr(self, e: TypedDictExpr) -> Type:
22482251
return self.expr_checker.visit_typeddict_expr(e)
22492252

‎mypy/checkexpr.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
2020
DictionaryComprehension, ComplexExpr, EllipsisExpr, StarExpr, AwaitExpr, YieldExpr,
2121
YieldFromExpr, TypedDictExpr, PromoteExpr, NewTypeExpr, NamedTupleExpr, TypeVarExpr,
22-
TypeAliasExpr, BackquoteExpr, ARG_POS, ARG_NAMED, ARG_STAR, ARG_STAR2, MODULE_REF,
23-
UNBOUND_TVAR, BOUND_TVAR,
22+
TypeAliasExpr, BackquoteExpr, EnumCallExpr,
23+
ARG_POS, ARG_NAMED, ARG_STAR, ARG_STAR2, MODULE_REF, UNBOUND_TVAR, BOUND_TVAR,
2424
)
2525
from mypy import nodes
2626
import mypy.checker
@@ -2156,6 +2156,10 @@ def visit_namedtuple_expr(self, e: NamedTupleExpr) -> Type:
21562156
# TODO: Perhaps return a type object type?
21572157
return AnyType()
21582158

2159+
def visit_enum_call_expr(self, e: EnumCallExpr) -> Type:
2160+
# TODO: Perhaps return a type object type?
2161+
return AnyType()
2162+
21592163
def visit_typeddict_expr(self, e: TypedDictExpr) -> Type:
21602164
# TODO: Perhaps return a type object type?
21612165
return AnyType()

‎mypy/nodes.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1792,6 +1792,22 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
17921792
return visitor.visit_typeddict_expr(self)
17931793

17941794

1795+
class EnumCallExpr(Expression):
1796+
"""Named tuple expression Enum('name', 'val1 val2 ...')."""
1797+
1798+
# The class representation of this enumerated type
1799+
info = None # type: TypeInfo
1800+
# The item names (for debugging)
1801+
items = None # type: List[str]
1802+
1803+
def __init__(self, info: 'TypeInfo', items: List[str]) -> None:
1804+
self.info = info
1805+
self.items = items
1806+
1807+
def accept(self, visitor: ExpressionVisitor[T]) -> T:
1808+
return visitor.visit_enum_call_expr(self)
1809+
1810+
17951811
class PromoteExpr(Expression):
17961812
"""Ducktype class decorator expression _promote(...)."""
17971813

‎mypy/semanal.py‎

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@
6464
YieldFromExpr, NamedTupleExpr, TypedDictExpr, NonlocalDecl, SymbolNode,
6565
SetComprehension, DictionaryComprehension, TYPE_ALIAS, TypeAliasExpr,
6666
YieldExpr, ExecStmt, Argument, BackquoteExpr, ImportBase, AwaitExpr,
67-
IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr, TempNode,
67+
IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr, TempNode, EnumCallExpr,
6868
COVARIANT, CONTRAVARIANT, INVARIANT, UNBOUND_IMPORTED, LITERAL_YES,
6969
)
7070
from mypy.typevars import has_no_typevars, fill_typevars
@@ -1195,6 +1195,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
11951195
self.process_typevar_declaration(s)
11961196
self.process_namedtuple_definition(s)
11971197
self.process_typeddict_definition(s)
1198+
self.process_enum_call(s)
11981199

11991200
if (len(s.lvalues) == 1 and isinstance(s.lvalues[0], NameExpr) and
12001201
s.lvalues[0].name == '__all__' and s.lvalues[0].kind == GDEF and
@@ -1991,6 +1992,103 @@ def build_typeddict_typeinfo(self, name: str, items: List[str],
19911992

19921993
return info
19931994

1995+
def process_enum_call(self, s: AssignmentStmt) -> None:
1996+
"""Check if s defines an Enum; if yes, store the definition in symbol table."""
1997+
if len(s.lvalues) != 1 or not isinstance(s.lvalues[0], NameExpr):
1998+
return
1999+
lvalue = s.lvalues[0]
2000+
name = lvalue.name
2001+
enum_call = self.check_enum_call(s.rvalue, name)
2002+
if enum_call is None:
2003+
return
2004+
# Yes, it's a valid Enum definition. Add it to the symbol table.
2005+
node = self.lookup(name, s)
2006+
if node:
2007+
node.kind = GDEF # TODO locally defined Enum
2008+
node.node = enum_call
2009+
2010+
def check_enum_call(self, node: Expression, var_name: str = None) -> Optional[TypeInfo]:
2011+
"""Check if a call defines an Enum.
2012+
2013+
Example:
2014+
2015+
A = enum.Enum('A', 'foo bar')
2016+
2017+
is equivalent to:
2018+
2019+
class A(enum.Enum):
2020+
foo = 1
2021+
bar = 2
2022+
"""
2023+
if not isinstance(node, CallExpr):
2024+
return None
2025+
call = node
2026+
callee = call.callee
2027+
if not isinstance(callee, RefExpr):
2028+
return None
2029+
fullname = callee.fullname
2030+
if fullname not in ('enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag'):
2031+
return None
2032+
items, types, ok = self.parse_enum_call_args(call)
2033+
if not ok:
2034+
# Error. Construct dummy return value.
2035+
return self.build_enum_call_typeinfo('Enum', [], [], fullname)
2036+
name = cast(StrExpr, call.args[0]).value
2037+
if name != var_name or self.is_func_scope():
2038+
# Give it a unique name derived from the line number.
2039+
name += '@' + str(call.line)
2040+
info = self.build_enum_call_typeinfo(name, items, types, fullname)
2041+
# Store it as a global just in case it would remain anonymous.
2042+
# (Or in the nearest class if there is one.)
2043+
stnode = SymbolTableNode(GDEF, info, self.cur_mod_id)
2044+
if self.type:
2045+
self.type.names[name] = stnode
2046+
else:
2047+
self.globals[name] = stnode
2048+
call.analyzed = EnumCallExpr(info, items)
2049+
call.analyzed.set_line(call.line, call.column)
2050+
return info
2051+
2052+
def build_enum_call_typeinfo(self, name: str, items: List[str],
2053+
types: List[Type], fullname: str) -> TypeInfo:
2054+
base = self.named_type_or_none(fullname)
2055+
assert base is not None
2056+
info = self.basic_new_typeinfo(name, base)
2057+
info.is_enum = True
2058+
for item, typ in zip(items, types):
2059+
var = Var(item, typ)
2060+
var.info = info
2061+
var.is_property = True
2062+
info.names[item] = SymbolTableNode(MDEF, var)
2063+
return info
2064+
2065+
def parse_enum_call_args(self, call: CallExpr) -> Tuple[List[str], List[Type], bool]:
2066+
args = call.args
2067+
if len(args) < 2:
2068+
return self.fail_enum_call_arg("Too few arguments for Enum()", call)
2069+
if len(args) > 2:
2070+
return self.fail_enum_call_arg("Too many arguments for Enum()", call)
2071+
if call.arg_kinds != [ARG_POS, ARG_POS]:
2072+
return self.fail_enum_call_arg("Unexpected arguments to Enum()", call)
2073+
if not isinstance(args[0], (StrExpr, BytesExpr, UnicodeExpr)):
2074+
return self.fail_typeddict_arg(
2075+
"Enum() expects a string literal as the first argument", call)
2076+
if not isinstance(args[1], (StrExpr, BytesExpr, UnicodeExpr)):
2077+
return self.fail_typeddict_arg(
2078+
"Enum() expects a string literal as the second argument", call)
2079+
fields = cast(StrExpr, args[1]).value
2080+
items = []
2081+
types = [] # type: List[Type]
2082+
for field in fields.replace(',', ' ').split():
2083+
items.append(field)
2084+
types.append(AnyType())
2085+
return items, types, True
2086+
2087+
def fail_enum_call_arg(self, message: str,
2088+
context: Context) -> Tuple[List[str], List[Type], bool]:
2089+
self.fail(message, context)
2090+
return [], [], False
2091+
19942092
def visit_decorator(self, dec: Decorator) -> None:
19952093
for d in dec.decorators:
19962094
d.accept(self)

‎mypy/strconv.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -429,6 +429,9 @@ def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> str:
429429
o.info.name(),
430430
o.info.tuple_type)
431431

432+
def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> str:
433+
return 'EnumCallExpr:{}({}, {})'.format(o.line, o.info.name(), o.items)
434+
432435
def visit_typeddict_expr(self, o: 'mypy.nodes.TypedDictExpr') -> str:
433436
return 'TypedDictExpr:{}({})'.format(o.line,
434437
o.info.name())

‎mypy/treetransform.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
ComparisonExpr, TempNode, StarExpr, Statement, Expression,
2020
YieldFromExpr, NamedTupleExpr, TypedDictExpr, NonlocalDecl, SetComprehension,
2121
DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
22-
YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr,
22+
YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr, EnumCallExpr,
2323
)
2424
from mypy.types import Type, FunctionLike
2525
from mypy.traverser import TraverserVisitor
@@ -483,6 +483,9 @@ def visit_newtype_expr(self, node: NewTypeExpr) -> NewTypeExpr:
483483
def visit_namedtuple_expr(self, node: NamedTupleExpr) -> NamedTupleExpr:
484484
return NamedTupleExpr(node.info)
485485

486+
def visit_enum_call_expr(self, node: EnumCallExpr) -> EnumCallExpr:
487+
return EnumCallExpr(node.info, list(node.items))
488+
486489
def visit_typeddict_expr(self, node: TypedDictExpr) -> Node:
487490
return TypedDictExpr(node.info)
488491

‎mypy/visitor.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,10 @@ def visit_type_alias_expr(self, o: 'mypy.nodes.TypeAliasExpr') -> T:
156156
def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
157157
pass
158158

159+
@abstractmethod
160+
def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> T:
161+
pass
162+
159163
@abstractmethod
160164
def visit_typeddict_expr(self, o: 'mypy.nodes.TypedDictExpr') -> T:
161165
pass
@@ -392,6 +396,9 @@ def visit_type_alias_expr(self, o: 'mypy.nodes.TypeAliasExpr') -> T:
392396
def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
393397
pass
394398

399+
def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> T:
400+
pass
401+
395402
def visit_typeddict_expr(self, o: 'mypy.nodes.TypedDictExpr') -> T:
396403
pass
397404

‎test-data/unit/pythoneval-enum.test‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,3 +132,17 @@ class E(N, Enum):
132132
def f(x: E) -> None: pass
133133

134134
f(E.X)
135+
136+
[case testFunctionalEnum]
137+
from enum import Enum, IntEnum
138+
E = Enum('E', 'foo bar')
139+
I = IntEnum('I', ' bar, baz ')
140+
reveal_type(E.foo)
141+
reveal_type(E.bar.value)
142+
reveal_type(I.bar)
143+
reveal_type(I.baz.value)
144+
[out]
145+
_program.py:4: error: Revealed type is '_testFunctionalEnum.E'
146+
_program.py:5: error: Revealed type is 'Any'
147+
_program.py:6: error: Revealed type is '_testFunctionalEnum.I'
148+
_program.py:7: error: Revealed type is 'builtins.int'

0 commit comments

Comments
 (0)