Skip to content

Commit 2d3dc1e

Browse files
Michael0x2agvanrossum
Michael0x2a
authored and
gvanrossum
committed
Add NewType (#1939)
This pull request implements NewType as described in PEP 484 and in issue #1284. It also adds a variety of test cases to verify NewType works correctly when used and misused and adds information about using NewType to the mypy docs. Fixes #1284.
1 parent 9ac8bf7 commit 2d3dc1e

11 files changed

Lines changed: 591 additions & 3 deletions

File tree

‎docs/source/kinds_of_types.rst‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,124 @@ A type alias does not create a new type. It's just a shorthand notation
425425
for another type -- it's equivalent to the target type. Type aliases
426426
can be imported from modules like any names.
427427

428+
.. _newtypes:
429+
430+
NewTypes
431+
********
432+
433+
(Freely after `PEP 484
434+
<https://www.python.org/dev/peps/pep-0484/#newtype-helper-function>`_.)
435+
436+
There are also situations where a programmer might want to avoid logical errors by
437+
creating simple classes. For example:
438+
439+
.. code-block:: python
440+
441+
class UserId(int):
442+
pass
443+
444+
get_by_user_id(user_id: UserId):
445+
...
446+
447+
However, this approach introduces some runtime overhead. To avoid this, the typing
448+
module provides a helper function ``NewType`` that creates simple unique types with
449+
almost zero runtime overhead. Mypy will treat the statement
450+
``Derived = NewType('Derived', Base)`` as being roughly equivalent to the following
451+
definition:
452+
453+
.. code-block:: python
454+
455+
class Derived(Base):
456+
def __init__(self, _x: Base) -> None:
457+
...
458+
459+
However, at runtime, ``NewType('Derived', Base)`` will return a dummy function that
460+
simply returns its argument:
461+
462+
.. code-block:: python
463+
464+
def Derived(_x):
465+
return _x
466+
467+
Mypy will require explicit casts from ``int`` where ``UserId`` is expected, while
468+
implicitly casting from ``UserId`` where ``int`` is expected. Examples:
469+
470+
.. code-block:: python
471+
472+
from typing import NewType
473+
474+
UserId = NewType('UserId', int)
475+
476+
def name_by_id(user_id: UserId) -> str:
477+
...
478+
479+
UserId('user') # Fails type check
480+
481+
name_by_id(42) # Fails type check
482+
name_by_id(UserId(42)) # OK
483+
484+
num = UserId(5) + 1 # type: int
485+
486+
``NewType`` accepts exactly two arguments. The first argument must be a string literal
487+
containing the name of the new type and must equal the name of the variable to which the new
488+
type is assigned. The second argument must be a properly subclassable class, i.e.,
489+
not a type construct like ``Union``, etc.
490+
491+
The function returned by ``NewType`` accepts only one argument; this is equivalent to
492+
supporting only one constructor accepting an instance of the base class (see above).
493+
Example:
494+
495+
.. code-block:: python
496+
497+
from typing import NewType
498+
499+
class PacketId:
500+
def __init__(self, major: int, minor: int) -> None:
501+
self._major = major
502+
self._minor = minor
503+
504+
TcpPacketId = NewType('TcpPacketId', PacketId)
505+
506+
packet = PacketId(100, 100)
507+
tcp_packet = TcpPacketId(packet) # OK
508+
509+
tcp_packet = TcpPacketId(127, 0) # Fails in type checker and at runtime
510+
511+
Both ``isinstance`` and ``issubclass``, as well as subclassing will fail for
512+
``NewType('Derived', Base)`` since function objects don't support these operations.
513+
514+
.. note::
515+
516+
Note that unlike type aliases, ``NewType`` will create an entirely new and
517+
unique type when used. The intended purpose of ``NewType`` is to help you
518+
detect cases where you accidentally mixed together the old base type and the
519+
new derived type.
520+
521+
For example, the following will successfully typecheck when using type
522+
aliases:
523+
524+
.. code-block:: python
525+
526+
UserId = int
527+
528+
def name_by_id(user_id: UserId) -> str:
529+
...
530+
531+
name_by_id(3) # ints and UserId are synonymous
532+
533+
But a similar example using ``NewType`` will not typecheck:
534+
535+
.. code-block:: python
536+
537+
from typing import NewType
538+
539+
UserId = NewType('UserId', int)
540+
541+
def name_by_id(user_id: UserId) -> str:
542+
...
543+
544+
name_by_id(3) # int is not the same as UserId
545+
428546
.. _named-tuples:
429547

430548
Named tuples

‎mypy/checker.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
BytesExpr, UnicodeExpr, FloatExpr, OpExpr, UnaryExpr, CastExpr, RevealTypeExpr, SuperExpr,
2121
TypeApplication, DictExpr, SliceExpr, FuncExpr, TempNode, SymbolTableNode,
2222
Context, ListComprehension, ConditionalExpr, GeneratorExpr,
23-
Decorator, SetExpr, TypeVarExpr, PrintStmt,
23+
Decorator, SetExpr, TypeVarExpr, NewTypeExpr, PrintStmt,
2424
LITERAL_TYPE, BreakStmt, ContinueStmt, ComparisonExpr, StarExpr,
2525
YieldFromExpr, NamedTupleExpr, SetComprehension,
2626
DictionaryComprehension, ComplexExpr, EllipsisExpr, TypeAliasExpr,
@@ -1979,6 +1979,9 @@ def visit_type_var_expr(self, e: TypeVarExpr) -> Type:
19791979
# TODO: Perhaps return a special type used for type variables only?
19801980
return AnyType()
19811981

1982+
def visit_newtype_expr(self, e: NewTypeExpr) -> Type:
1983+
return AnyType()
1984+
19821985
def visit_namedtuple_expr(self, e: NamedTupleExpr) -> Type:
19831986
# TODO: Perhaps return a type object type?
19841987
return AnyType()

‎mypy/nodes.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1710,6 +1710,18 @@ def accept(self, visitor: NodeVisitor[T]) -> T:
17101710
return visitor.visit__promote_expr(self)
17111711

17121712

1713+
class NewTypeExpr(Expression):
1714+
"""NewType expression NewType(...)."""
1715+
1716+
info = None # type: Optional[TypeInfo]
1717+
1718+
def __init__(self, info: Optional['TypeInfo']) -> None:
1719+
self.info = info
1720+
1721+
def accept(self, visitor: NodeVisitor[T]) -> T:
1722+
return visitor.visit_newtype_expr(self)
1723+
1724+
17131725
class AwaitExpr(Node):
17141726
"""Await expression (await ...)."""
17151727

@@ -1798,6 +1810,9 @@ class is generic then it will be a type constructor of higher kind.
17981810
# Is this a named tuple type?
17991811
is_named_tuple = False
18001812

1813+
# Is this a newtype type?
1814+
is_newtype = False
1815+
18011816
# Is this a dummy from deserialization?
18021817
is_dummy = False
18031818

@@ -1974,6 +1989,7 @@ def serialize(self) -> Union[str, JsonDict]:
19741989
'_promote': None if self._promote is None else self._promote.serialize(),
19751990
'tuple_type': None if self.tuple_type is None else self.tuple_type.serialize(),
19761991
'is_named_tuple': self.is_named_tuple,
1992+
'is_newtype': self.is_newtype,
19771993
}
19781994
return data
19791995

@@ -1996,6 +2012,7 @@ def deserialize(cls, data: JsonDict) -> 'TypeInfo':
19962012
ti.tuple_type = (None if data['tuple_type'] is None
19972013
else mypy.types.TupleType.deserialize(data['tuple_type']))
19982014
ti.is_named_tuple = data['is_named_tuple']
2015+
ti.is_newtype = data['is_newtype']
19992016
return ti
20002017

20012018

‎mypy/semanal.py‎

Lines changed: 116 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@
5858
GlobalDecl, SuperExpr, DictExpr, CallExpr, RefExpr, OpExpr, UnaryExpr,
5959
SliceExpr, CastExpr, RevealTypeExpr, TypeApplication, Context, SymbolTable,
6060
SymbolTableNode, BOUND_TVAR, UNBOUND_TVAR, ListComprehension, GeneratorExpr,
61-
FuncExpr, MDEF, FuncBase, Decorator, SetExpr, TypeVarExpr,
61+
FuncExpr, MDEF, FuncBase, Decorator, SetExpr, TypeVarExpr, NewTypeExpr,
6262
StrExpr, BytesExpr, PrintStmt, ConditionalExpr, PromoteExpr,
6363
ComparisonExpr, StarExpr, ARG_POS, ARG_NAMED, MroError, type_aliases,
6464
YieldFromExpr, NamedTupleExpr, NonlocalDecl,
@@ -769,6 +769,8 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
769769
defn.info.tuple_type = base
770770
base_types.append(base.fallback)
771771
elif isinstance(base, Instance):
772+
if base.type.is_newtype:
773+
self.fail("Cannot subclass NewType", defn)
772774
base_types.append(base)
773775
elif isinstance(base, AnyType):
774776
defn.info.fallback_to_any = True
@@ -1083,6 +1085,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
10831085
for lvalue in s.lvalues:
10841086
self.store_declared_types(lvalue, s.type)
10851087
self.check_and_set_up_type_alias(s)
1088+
self.process_newtype_declaration(s)
10861089
self.process_typevar_declaration(s)
10871090
self.process_namedtuple_definition(s)
10881091

@@ -1289,6 +1292,118 @@ def store_declared_types(self, lvalue: Node, typ: Type) -> None:
12891292
# This has been flagged elsewhere as an error, so just ignore here.
12901293
pass
12911294

1295+
def process_newtype_declaration(self, s: AssignmentStmt) -> None:
1296+
"""Check if s declares a NewType; if yes, store it in symbol table."""
1297+
# Extract and check all information from newtype declaration
1298+
name, call = self.analyze_newtype_declaration(s)
1299+
if name is None or call is None:
1300+
return
1301+
1302+
old_type = self.check_newtype_args(name, call, s)
1303+
if old_type is None:
1304+
return
1305+
1306+
# Create the corresponding class definition if the aliased type is subtypeable
1307+
if isinstance(old_type, TupleType):
1308+
newtype_class_info = self.build_newtype_typeinfo(name, old_type, old_type.fallback)
1309+
newtype_class_info.tuple_type = old_type
1310+
elif isinstance(old_type, Instance):
1311+
newtype_class_info = self.build_newtype_typeinfo(name, old_type, old_type)
1312+
else:
1313+
message = "Argument 2 to NewType(...) must be subclassable (got {})"
1314+
self.fail(message.format(old_type), s)
1315+
return
1316+
1317+
# If so, add it to the symbol table.
1318+
node = self.lookup(name, s)
1319+
if node is None:
1320+
self.fail("Could not find {} in current namespace".format(name), s)
1321+
return
1322+
# TODO: why does NewType work in local scopes despite always being of kind GDEF?
1323+
node.kind = GDEF
1324+
node.node = newtype_class_info
1325+
call.analyzed = NewTypeExpr(newtype_class_info).set_line(call.line)
1326+
1327+
def analyze_newtype_declaration(self,
1328+
s: AssignmentStmt) -> Tuple[Optional[str], Optional[CallExpr]]:
1329+
"""Return the NewType call expression if `s` is a newtype declaration or None otherwise."""
1330+
name, call = None, None
1331+
if (len(s.lvalues) == 1
1332+
and isinstance(s.lvalues[0], NameExpr)
1333+
and isinstance(s.rvalue, CallExpr)
1334+
and isinstance(s.rvalue.callee, RefExpr)
1335+
and s.rvalue.callee.fullname == 'typing.NewType'):
1336+
lvalue = s.lvalues[0]
1337+
name = s.lvalues[0].name
1338+
if not lvalue.is_def:
1339+
if s.type:
1340+
self.fail("Cannot declare the type of a NewType declaration", s)
1341+
else:
1342+
self.fail("Cannot redefine '%s' as a NewType" % name, s)
1343+
1344+
# This dummy NewTypeExpr marks the call as sufficiently analyzed; it will be
1345+
# overwritten later with a fully complete NewTypeExpr if there are no other
1346+
# errors with the NewType() call.
1347+
call = s.rvalue
1348+
call.analyzed = NewTypeExpr(None).set_line(call.line)
1349+
1350+
return name, call
1351+
1352+
def check_newtype_args(self, name: str, call: CallExpr, context: Context) -> Optional[Type]:
1353+
has_failed = False
1354+
args, arg_kinds = call.args, call.arg_kinds
1355+
if len(args) != 2 or arg_kinds[0] != ARG_POS or arg_kinds[1] != ARG_POS:
1356+
self.fail("NewType(...) expects exactly two positional arguments", context)
1357+
return None
1358+
1359+
# Check first argument
1360+
if not isinstance(args[0], (StrExpr, BytesExpr, UnicodeExpr)):
1361+
self.fail("Argument 1 to NewType(...) must be a string literal", context)
1362+
has_failed = True
1363+
elif cast(StrExpr, call.args[0]).value != name:
1364+
self.fail("Argument 1 to NewType(...) does not match variable name", context)
1365+
has_failed = True
1366+
1367+
# Check second argument
1368+
try:
1369+
unanalyzed_type = expr_to_unanalyzed_type(call.args[1])
1370+
except TypeTranslationError:
1371+
self.fail("Argument 2 to NewType(...) must be a valid type", context)
1372+
return None
1373+
old_type = self.anal_type(unanalyzed_type)
1374+
1375+
if isinstance(old_type, Instance) and old_type.type.is_newtype:
1376+
self.fail("Argument 2 to NewType(...) cannot be another NewType", context)
1377+
has_failed = True
1378+
1379+
return None if has_failed else old_type
1380+
1381+
def build_newtype_typeinfo(self, name: str, old_type: Type, base_type: Instance) -> TypeInfo:
1382+
class_def = ClassDef(name, Block([]))
1383+
class_def.fullname = self.qualified_name(name)
1384+
1385+
symbols = SymbolTable()
1386+
info = TypeInfo(symbols, class_def)
1387+
info.mro = [info] + base_type.type.mro
1388+
info.bases = [base_type]
1389+
info.is_newtype = True
1390+
1391+
# Add __init__ method
1392+
args = [Argument(Var('cls'), NoneTyp(), None, ARG_POS),
1393+
self.make_argument('item', old_type)]
1394+
signature = CallableType(
1395+
arg_types=[cast(Type, None), old_type],
1396+
arg_kinds=[arg.kind for arg in args],
1397+
arg_names=['self', 'item'],
1398+
ret_type=old_type,
1399+
fallback=self.named_type('__builtins__.function'),
1400+
name=name)
1401+
init_func = FuncDef('__init__', args, Block([]), typ=signature)
1402+
init_func.info = info
1403+
symbols['__init__'] = SymbolTableNode(MDEF, init_func)
1404+
1405+
return info
1406+
12921407
def process_typevar_declaration(self, s: AssignmentStmt) -> None:
12931408
"""Check if s declares a TypeVar; it yes, store it in symbol table."""
12941409
call = self.get_typevar_declaration(s)

‎mypy/strconv.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,9 @@ def visit_namedtuple_expr(self, o):
433433
def visit__promote_expr(self, o):
434434
return 'PromoteExpr:{}({})'.format(o.line, o.type)
435435

436+
def visit_newtype_expr(self, o):
437+
return 'NewTypeExpr:{}({}, {})'.format(o.line, o.fullname(), self.dump([o.value], o))
438+
436439
def visit_func_expr(self, o):
437440
a = self.func_helper(o)
438441
return self.dump(a, o)

‎mypy/test/testcheck.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@
6363
'check-fastparse.test',
6464
'check-warnings.test',
6565
'check-async-await.test',
66+
'check-newtype.test',
6667
]
6768

6869

‎mypy/treetransform.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
ConditionalExpr, DictExpr, SetExpr, NameExpr, IntExpr, StrExpr, BytesExpr,
1616
UnicodeExpr, FloatExpr, CallExpr, SuperExpr, MemberExpr, IndexExpr,
1717
SliceExpr, OpExpr, UnaryExpr, FuncExpr, TypeApplication, PrintStmt,
18-
SymbolTable, RefExpr, TypeVarExpr, PromoteExpr,
18+
SymbolTable, RefExpr, TypeVarExpr, NewTypeExpr, PromoteExpr,
1919
ComparisonExpr, TempNode, StarExpr,
2020
YieldFromExpr, NamedTupleExpr, NonlocalDecl, SetComprehension,
2121
DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
@@ -453,6 +453,9 @@ def visit_type_var_expr(self, node: TypeVarExpr) -> Node:
453453
def visit_type_alias_expr(self, node: TypeAliasExpr) -> TypeAliasExpr:
454454
return TypeAliasExpr(node.type)
455455

456+
def visit_newtype_expr(self, node: NewTypeExpr) -> NewTypeExpr:
457+
return NewTypeExpr(node.info)
458+
456459
def visit_namedtuple_expr(self, node: NamedTupleExpr) -> Node:
457460
return NamedTupleExpr(node.info)
458461

‎mypy/visitor.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,9 @@ def visit_type_alias_expr(self, o: 'mypy.nodes.TypeAliasExpr') -> T:
225225
def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
226226
pass
227227

228+
def visit_newtype_expr(self, o: 'mypy.nodes.NewTypeExpr') -> T:
229+
pass
230+
228231
def visit__promote_expr(self, o: 'mypy.nodes.PromoteExpr') -> T:
229232
pass
230233

0 commit comments

Comments
 (0)