Support NewType in stubs.

This started as an attempt to support the stubs added to the pytype exclude list in python/typeshed#4785. Those stubs turned out to be a bit of a hairball, but supporting NewType seems useful anyway.

Resolves #597.

PiperOrigin-RevId: 345364292

Author: rchen152

pytype/pyi/lexer.lex

@@ -103,6 +103,8 @@ typedef pytype::parser::token t;
 "typing.NamedTuple" { return t::NAMEDTUPLE; }
 "namedtuple" { return t::COLL_NAMEDTUPLE; }
 "collections.namedtuple" { return t::COLL_NAMEDTUPLE; }
+"NewType" { return t::NEWTYPE; }
+"typing.NewType" { return t::NEWTYPE; }
 "TypedDict" { return t::TYPEDDICT; }
 "typing.TypedDict" { return t::TYPEDDICT; }
 "typing_extensions.TypedDict" { return t::TYPEDDICT; }

pytype/pyi/lexer_test.py

@@ -302,6 +302,10 @@ def test_typeddict(self):
     self.check(["TYPEDDICT"], "typing.TypedDict")
     self.check(["TYPEDDICT"], "typing_extensions.TypedDict")
 
+  def test_newtype(self):
+    self.check(["NEWTYPE"], "NewType")
+    self.check(["NEWTYPE"], "typing.NewType")
+
 
 if __name__ == "__main__":
   unittest.main()

pytype/pyi/parser.h

@@ -34,6 +34,7 @@ enum CallSelector {
   kNewConstant,
   kNewFunction,
   kNewNamedTuple,
+  kNewNewType,
   kNewTypedDict,
   kRegisterClassName,
   kAddTypeVar,

pytype/pyi/parser.py

@@ -1081,6 +1081,27 @@ def new_named_tuple(self, base_name, fields):
     self._generated_classes[base_name].append(nt_class)
     return pytd.NamedType(nt_class.name)
 
+  def new_new_type(self, name, typ):
+    """Returns a type for a NewType."""
+    name = _handle_string_literal(name)
+    args = [("self", pytd.AnythingType(), None), ("val", typ, None)]
+    ret = pytd.NamedType("NoneType")
+    methods = _merge_method_signatures(
+        [self.new_function((), False, "__init__", args, ret, ())])
+    cls_name = escape.pack_newtype_base_class(
+        name, len(self._generated_classes[name]))
+    cls = pytd.Class(name=cls_name,
+                     metaclass=None,
+                     parents=(typ,),
+                     methods=tuple(methods),
+                     constants=(),
+                     decorators=(),
+                     classes=(),
+                     slots=None,
+                     template=())
+    self._generated_classes[name].append(cls)
+    return pytd.NamedType(cls_name)
+
   def new_typed_dict(self, name, items, total):
     """Returns a type for a TypedDict.
 

pytype/pyi/parser.yy

@@ -70,7 +70,7 @@ int pytypelex(pytype::parser::semantic_type* lvalp, pytype::location* llocp,
 
 /* Reserved words. */
 %token ASYNC CLASS DEF ELSE ELIF IF OR AND PASS IMPORT FROM AS RAISE
-%token NOTHING NAMEDTUPLE COLL_NAMEDTUPLE TYPEDDICT TYPEVAR
+%token NOTHING NAMEDTUPLE COLL_NAMEDTUPLE NEWTYPE TYPEDDICT TYPEVAR
 /* Punctuation. */
 %token ARROW ELLIPSIS EQ NE LE GE
 /* Other. */
@@ -435,6 +435,9 @@ from_item
   | COLL_NAMEDTUPLE {
       $$ = PyString_FromString("namedtuple");
     }
+  | NEWTYPE {
+      $$ = PyString_FromString("NewType");
+    }
   | TYPEDDICT {
       $$ = PyString_FromString("TypedDict");
     }
@@ -445,6 +448,7 @@ from_item
       $$ = PyString_FromString("*");
     }
   | NAME AS NAME { $$ = Py_BuildValue("(NN)", $1, $3); }
+  | NEWTYPE AS NEWTYPE { $$ = PyString_FromString("NewType"); }
   ;
 
 alias_or_constant
@@ -504,6 +508,7 @@ funcdef
 funcname
   : NAME { $$ = $1; }
   | COLL_NAMEDTUPLE { $$ = PyString_FromString("namedtuple"); }
+  | NEWTYPE { $$ = PyString_FromString("NewType"); }
   | TYPEDDICT { $$ = PyString_FromString("TypedDict"); }
   ;
 
@@ -652,6 +657,10 @@ type
       $$ = ctx->Call(kNewNamedTuple, "(NN)", $3, $5);
       CHECK($$, @$);
     }
+  | NEWTYPE '(' STRING ',' type maybe_comma ')' {
+      $$ = ctx->Call(kNewNewType, "(NN)", $3, $5);
+      CHECK($$, @$);
+    }
   | TYPEDDICT '(' STRING ',' typed_dict_fields maybe_typed_dict_kwarg ')' {
       $$ = ctx->Call(kNewTypedDict, "(NNN)", $3, $5, $6);
       CHECK($$, @$);

pytype/pyi/parser_ext.cc

@@ -45,6 +45,7 @@ static const SelectorEntry<CallSelector> call_attributes[] = {
   {kNewConstant, "new_constant"},
   {kNewFunction, "new_function"},
   {kNewNamedTuple, "new_named_tuple"},
+  {kNewNewType, "new_new_type"},
   {kNewTypedDict, "new_typed_dict"},
   {kRegisterClassName, "register_class_name"},
   {kAddTypeVar, "add_type_var"},
@@ -294,6 +295,7 @@ static void add_tokens_dict(PyObject* module) {
   add_token(tokens, "NOTHING",         t::NOTHING);
   add_token(tokens, "NAMEDTUPLE",      t::NAMEDTUPLE);
   add_token(tokens, "COLL_NAMEDTUPLE", t::COLL_NAMEDTUPLE);
+  add_token(tokens, "NEWTYPE",         t::NEWTYPE);
   add_token(tokens, "TYPEDDICT",       t::TYPEDDICT);
   add_token(tokens, "TYPEVAR",         t::TYPEVAR);
 

pytype/pyi/parser_test.py

@@ -2579,5 +2579,30 @@ class Foo(TypedDict, metaclass=Meta): ...
     """)
 
 
+class NewTypeTest(_ParserTestBase):
+
+  def test_basic(self):
+    self.check("""
+      from typing import NewType
+      X = NewType('X', int)
+    """, """
+      X = newtype_X_0
+
+      class newtype_X_0(int):
+          def __init__(self, val: int) -> None: ...
+    """)
+
+  def test_fullname(self):
+    self.check("""
+      import typing
+      X = typing.NewType('X', int)
+    """, """
+      X = newtype_X_0
+
+      class newtype_X_0(int):
+          def __init__(self, val: int) -> None: ...
+    """)
+
+
 if __name__ == "__main__":
   unittest.main()

pytype/pytd/escape.py

@@ -61,3 +61,8 @@ def pack_namedtuple_base_class(name: str, index: int) -> str:
 def unpack_namedtuple(name: str) -> str:
   """Retrieve the original namedtuple class name."""
   return re.sub(r"\bnamedtuple[-_]([^-_]+)[-_\w]*", r"\1", name)
+
+
+def pack_newtype_base_class(name: str, index: int) -> str:
+  """Generate a name for a NewType proxy base class."""
+  return "newtype_%s_%d" % (name, index)