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.

Author: Michael0x2a

docs/source/kinds_of_types.rst

@@ -425,6 +425,124 @@ A type alias does not create a new type. It's just a shorthand notation
 for another type -- it's equivalent to the target type. Type aliases
 can be imported from modules like any names.
 
+.. _newtypes:
+
+NewTypes
+********
+
+(Freely after `PEP 484
+<https://www.python.org/dev/peps/pep-0484/#newtype-helper-function>`_.)
+
+There are also situations where a programmer might want to avoid logical errors by
+creating simple classes. For example:
+
+.. code-block:: python
+
+    class UserId(int):
+        pass
+
+    get_by_user_id(user_id: UserId):
+        ...
+
+However, this approach introduces some runtime overhead. To avoid this, the typing
+module provides a helper function ``NewType`` that creates simple unique types with
+almost zero runtime overhead. Mypy will treat the statement
+``Derived = NewType('Derived', Base)`` as being roughly equivalent to the following
+definition:
+
+.. code-block:: python
+
+    class Derived(Base):
+        def __init__(self, _x: Base) -> None:
+            ...
+
+However, at runtime, ``NewType('Derived', Base)`` will return a dummy function that
+simply returns its argument:
+
+.. code-block:: python
+
+    def Derived(_x):
+        return _x
+
+Mypy will require explicit casts from ``int`` where ``UserId`` is expected, while
+implicitly casting from ``UserId`` where ``int`` is expected. Examples:
+
+.. code-block:: python
+
+    from typing import NewType
+
+    UserId = NewType('UserId', int)
+
+    def name_by_id(user_id: UserId) -> str:
+        ...
+
+    UserId('user')          # Fails type check
+
+    name_by_id(42)          # Fails type check
+    name_by_id(UserId(42))  # OK
+
+    num = UserId(5) + 1     # type: int
+
+``NewType`` accepts exactly two arguments. The first argument must be a string literal
+containing the name of the new type and must equal the name of the variable to which the new
+type is assigned. The second argument must be a properly subclassable class, i.e.,
+not a type construct like ``Union``, etc.
+
+The function returned by ``NewType`` accepts only one argument; this is equivalent to
+supporting only one constructor accepting an instance of the base class (see above).
+Example:
+
+.. code-block:: python
+
+    from typing import NewType
+
+    class PacketId:
+        def __init__(self, major: int, minor: int) -> None:
+            self._major = major
+            self._minor = minor
+
+    TcpPacketId = NewType('TcpPacketId', PacketId)
+
+    packet = PacketId(100, 100)
+    tcp_packet = TcpPacketId(packet)  # OK
+
+    tcp_packet = TcpPacketId(127, 0)  # Fails in type checker and at runtime
+
+Both ``isinstance`` and ``issubclass``, as well as subclassing will fail for
+``NewType('Derived', Base)`` since function objects don't support these operations.
+
+.. note::
+
+    Note that unlike type aliases, ``NewType`` will create an entirely new and
+    unique type when used. The intended purpose of ``NewType`` is to help you
+    detect cases where you accidentally mixed together the old base type and the
+    new derived type.
+
+    For example, the following will successfully typecheck when using type
+    aliases:
+
+    .. code-block:: python
+
+        UserId = int
+
+        def name_by_id(user_id: UserId) -> str:
+            ...
+
+        name_by_id(3)  # ints and UserId are synonymous
+
+    But a similar example using ``NewType`` will not typecheck:
+
+    .. code-block:: python
+
+        from typing import NewType
+
+        UserId = NewType('UserId', int)
+
+        def name_by_id(user_id: UserId) -> str:
+            ...
+
+        name_by_id(3)  # int is not the same as UserId
+
 .. _named-tuples:
 
 Named tuples

mypy/checker.py

@@ -20,7 +20,7 @@
     BytesExpr, UnicodeExpr, FloatExpr, OpExpr, UnaryExpr, CastExpr, RevealTypeExpr, SuperExpr,
     TypeApplication, DictExpr, SliceExpr, FuncExpr, TempNode, SymbolTableNode,
     Context, ListComprehension, ConditionalExpr, GeneratorExpr,
-    Decorator, SetExpr, TypeVarExpr, PrintStmt,
+    Decorator, SetExpr, TypeVarExpr, NewTypeExpr, PrintStmt,
     LITERAL_TYPE, BreakStmt, ContinueStmt, ComparisonExpr, StarExpr,
     YieldFromExpr, NamedTupleExpr, SetComprehension,
     DictionaryComprehension, ComplexExpr, EllipsisExpr, TypeAliasExpr,
@@ -1979,6 +1979,9 @@ def visit_type_var_expr(self, e: TypeVarExpr) -> Type:
         # TODO: Perhaps return a special type used for type variables only?
         return AnyType()
 
+    def visit_newtype_expr(self, e: NewTypeExpr) -> Type:
+        return AnyType()
+
     def visit_namedtuple_expr(self, e: NamedTupleExpr) -> Type:
         # TODO: Perhaps return a type object type?
         return AnyType()

mypy/nodes.py

@@ -1710,6 +1710,18 @@ def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit__promote_expr(self)
 
 
+class NewTypeExpr(Expression):
+    """NewType expression NewType(...)."""
+
+    info = None  # type: Optional[TypeInfo]
+
+    def __init__(self, info: Optional['TypeInfo']) -> None:
+        self.info = info
+
+    def accept(self, visitor: NodeVisitor[T]) -> T:
+        return visitor.visit_newtype_expr(self)
+
+
 class AwaitExpr(Node):
     """Await expression (await ...)."""
 
@@ -1798,6 +1810,9 @@ class is generic then it will be a type constructor of higher kind.
     # Is this a named tuple type?
     is_named_tuple = False
 
+    # Is this a newtype type?
+    is_newtype = False
+
     # Is this a dummy from deserialization?
     is_dummy = False
 
@@ -1974,6 +1989,7 @@ def serialize(self) -> Union[str, JsonDict]:
                 '_promote': None if self._promote is None else self._promote.serialize(),
                 'tuple_type': None if self.tuple_type is None else self.tuple_type.serialize(),
                 'is_named_tuple': self.is_named_tuple,
+                'is_newtype': self.is_newtype,
                 }
         return data
 
@@ -1996,6 +2012,7 @@ def deserialize(cls, data: JsonDict) -> 'TypeInfo':
         ti.tuple_type = (None if data['tuple_type'] is None
                          else mypy.types.TupleType.deserialize(data['tuple_type']))
         ti.is_named_tuple = data['is_named_tuple']
+        ti.is_newtype = data['is_newtype']
         return ti
 
 

mypy/semanal.py

@@ -58,7 +58,7 @@
     GlobalDecl, SuperExpr, DictExpr, CallExpr, RefExpr, OpExpr, UnaryExpr,
     SliceExpr, CastExpr, RevealTypeExpr, TypeApplication, Context, SymbolTable,
     SymbolTableNode, BOUND_TVAR, UNBOUND_TVAR, ListComprehension, GeneratorExpr,
-    FuncExpr, MDEF, FuncBase, Decorator, SetExpr, TypeVarExpr,
+    FuncExpr, MDEF, FuncBase, Decorator, SetExpr, TypeVarExpr, NewTypeExpr,
     StrExpr, BytesExpr, PrintStmt, ConditionalExpr, PromoteExpr,
     ComparisonExpr, StarExpr, ARG_POS, ARG_NAMED, MroError, type_aliases,
     YieldFromExpr, NamedTupleExpr, NonlocalDecl,
@@ -769,6 +769,8 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
                 defn.info.tuple_type = base
                 base_types.append(base.fallback)
             elif isinstance(base, Instance):
+                if base.type.is_newtype:
+                    self.fail("Cannot subclass NewType", defn)
                 base_types.append(base)
             elif isinstance(base, AnyType):
                 defn.info.fallback_to_any = True
@@ -1083,6 +1085,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
             for lvalue in s.lvalues:
                 self.store_declared_types(lvalue, s.type)
         self.check_and_set_up_type_alias(s)
+        self.process_newtype_declaration(s)
         self.process_typevar_declaration(s)
         self.process_namedtuple_definition(s)
 
@@ -1289,6 +1292,118 @@ def store_declared_types(self, lvalue: Node, typ: Type) -> None:
             # This has been flagged elsewhere as an error, so just ignore here.
             pass
 
+    def process_newtype_declaration(self, s: AssignmentStmt) -> None:
+        """Check if s declares a NewType; if yes, store it in symbol table."""
+        # Extract and check all information from newtype declaration
+        name, call = self.analyze_newtype_declaration(s)
+        if name is None or call is None:
+            return
+
+        old_type = self.check_newtype_args(name, call, s)
+        if old_type is None:
+            return
+
+        # Create the corresponding class definition if the aliased type is subtypeable
+        if isinstance(old_type, TupleType):
+            newtype_class_info = self.build_newtype_typeinfo(name, old_type, old_type.fallback)
+            newtype_class_info.tuple_type = old_type
+        elif isinstance(old_type, Instance):
+            newtype_class_info = self.build_newtype_typeinfo(name, old_type, old_type)
+        else:
+            message = "Argument 2 to NewType(...) must be subclassable (got {})"
+            self.fail(message.format(old_type), s)
+            return
+
+        # If so, add it to the symbol table.
+        node = self.lookup(name, s)
+        if node is None:
+            self.fail("Could not find {} in current namespace".format(name), s)
+            return
+        # TODO: why does NewType work in local scopes despite always being of kind GDEF?
+        node.kind = GDEF
+        node.node = newtype_class_info
+        call.analyzed = NewTypeExpr(newtype_class_info).set_line(call.line)
+
+    def analyze_newtype_declaration(self,
+            s: AssignmentStmt) -> Tuple[Optional[str], Optional[CallExpr]]:
+        """Return the NewType call expression if `s` is a newtype declaration or None otherwise."""
+        name, call = None, None
+        if (len(s.lvalues) == 1
+                and isinstance(s.lvalues[0], NameExpr)
+                and isinstance(s.rvalue, CallExpr)
+                and isinstance(s.rvalue.callee, RefExpr)
+                and s.rvalue.callee.fullname == 'typing.NewType'):
+            lvalue = s.lvalues[0]
+            name = s.lvalues[0].name
+            if not lvalue.is_def:
+                if s.type:
+                    self.fail("Cannot declare the type of a NewType declaration", s)
+                else:
+                    self.fail("Cannot redefine '%s' as a NewType" % name, s)
+
+            # This dummy NewTypeExpr marks the call as sufficiently analyzed; it will be
+            # overwritten later with a fully complete NewTypeExpr if there are no other
+            # errors with the NewType() call.
+            call = s.rvalue
+            call.analyzed = NewTypeExpr(None).set_line(call.line)
+
+        return name, call
+
+    def check_newtype_args(self, name: str, call: CallExpr, context: Context) -> Optional[Type]:
+        has_failed = False
+        args, arg_kinds = call.args, call.arg_kinds
+        if len(args) != 2 or arg_kinds[0] != ARG_POS or arg_kinds[1] != ARG_POS:
+            self.fail("NewType(...) expects exactly two positional arguments", context)
+            return None
+
+        # Check first argument
+        if not isinstance(args[0], (StrExpr, BytesExpr, UnicodeExpr)):
+            self.fail("Argument 1 to NewType(...) must be a string literal", context)
+            has_failed = True
+        elif cast(StrExpr, call.args[0]).value != name:
+            self.fail("Argument 1 to NewType(...) does not match variable name", context)
+            has_failed = True
+
+        # Check second argument
+        try:
+            unanalyzed_type = expr_to_unanalyzed_type(call.args[1])
+        except TypeTranslationError:
+            self.fail("Argument 2 to NewType(...) must be a valid type", context)
+            return None
+        old_type = self.anal_type(unanalyzed_type)
+
+        if isinstance(old_type, Instance) and old_type.type.is_newtype:
+            self.fail("Argument 2 to NewType(...) cannot be another NewType", context)
+            has_failed = True
+
+        return None if has_failed else old_type
+
+    def build_newtype_typeinfo(self, name: str, old_type: Type, base_type: Instance) -> TypeInfo:
+        class_def = ClassDef(name, Block([]))
+        class_def.fullname = self.qualified_name(name)
+
+        symbols = SymbolTable()
+        info = TypeInfo(symbols, class_def)
+        info.mro = [info] + base_type.type.mro
+        info.bases = [base_type]
+        info.is_newtype = True
+
+        # Add __init__ method
+        args = [Argument(Var('cls'), NoneTyp(), None, ARG_POS),
+                self.make_argument('item', old_type)]
+        signature = CallableType(
+            arg_types=[cast(Type, None), old_type],
+            arg_kinds=[arg.kind for arg in args],
+            arg_names=['self', 'item'],
+            ret_type=old_type,
+            fallback=self.named_type('__builtins__.function'),
+            name=name)
+        init_func = FuncDef('__init__', args, Block([]), typ=signature)
+        init_func.info = info
+        symbols['__init__'] = SymbolTableNode(MDEF, init_func)
+
+        return info
+
     def process_typevar_declaration(self, s: AssignmentStmt) -> None:
         """Check if s declares a TypeVar; it yes, store it in symbol table."""
         call = self.get_typevar_declaration(s)

mypy/strconv.py

@@ -433,6 +433,9 @@ def visit_namedtuple_expr(self, o):
     def visit__promote_expr(self, o):
         return 'PromoteExpr:{}({})'.format(o.line, o.type)
 
+    def visit_newtype_expr(self, o):
+        return 'NewTypeExpr:{}({}, {})'.format(o.line, o.fullname(), self.dump([o.value], o))
+
     def visit_func_expr(self, o):
         a = self.func_helper(o)
         return self.dump(a, o)

mypy/test/testcheck.py

@@ -63,6 +63,7 @@
     'check-fastparse.test',
     'check-warnings.test',
     'check-async-await.test',
+    'check-newtype.test',
 ]
 
 

mypy/treetransform.py

@@ -15,7 +15,7 @@
     ConditionalExpr, DictExpr, SetExpr, NameExpr, IntExpr, StrExpr, BytesExpr,
     UnicodeExpr, FloatExpr, CallExpr, SuperExpr, MemberExpr, IndexExpr,
     SliceExpr, OpExpr, UnaryExpr, FuncExpr, TypeApplication, PrintStmt,
-    SymbolTable, RefExpr, TypeVarExpr, PromoteExpr,
+    SymbolTable, RefExpr, TypeVarExpr, NewTypeExpr, PromoteExpr,
     ComparisonExpr, TempNode, StarExpr,
     YieldFromExpr, NamedTupleExpr, NonlocalDecl, SetComprehension,
     DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
@@ -453,6 +453,9 @@ def visit_type_var_expr(self, node: TypeVarExpr) -> Node:
     def visit_type_alias_expr(self, node: TypeAliasExpr) -> TypeAliasExpr:
         return TypeAliasExpr(node.type)
 
+    def visit_newtype_expr(self, node: NewTypeExpr) -> NewTypeExpr:
+        return NewTypeExpr(node.info)
+
     def visit_namedtuple_expr(self, node: NamedTupleExpr) -> Node:
         return NamedTupleExpr(node.info)
 

mypy/visitor.py

@@ -225,6 +225,9 @@ def visit_type_alias_expr(self, o: 'mypy.nodes.TypeAliasExpr') -> T:
     def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
         pass
 
+    def visit_newtype_expr(self, o: 'mypy.nodes.NewTypeExpr') -> T:
+        pass
+
     def visit__promote_expr(self, o: 'mypy.nodes.PromoteExpr') -> T:
         pass