New backward-compatible NamedTuple (#282)

Author: ilevkivskyi

src/test_typing.py

@@ -1169,6 +1169,10 @@ class CSub(B):
     z: ClassVar['CSub'] = B()
 class G(Generic[T]):
     lst: ClassVar[List[T]] = []
+
+class CoolEmployee(NamedTuple):
+    name: str
+    cool: int
 """
 
 if PY36:
@@ -1586,6 +1590,17 @@ def test_basics(self):
         self.assertEqual(Emp._fields, ('name', 'id'))
         self.assertEqual(Emp._field_types, dict(name=str, id=int))
 
+    @skipUnless(PY36, 'Python 3.6 required')
+    def test_annotation_usage(self):
+        tim = CoolEmployee('Tim', 9000)
+        self.assertIsInstance(tim, CoolEmployee)
+        self.assertIsInstance(tim, tuple)
+        self.assertEqual(tim.name, 'Tim')
+        self.assertEqual(tim.cool, 9000)
+        self.assertEqual(CoolEmployee.__name__, 'CoolEmployee')
+        self.assertEqual(CoolEmployee._fields, ('name', 'cool'))
+        self.assertEqual(CoolEmployee._field_types, dict(name=str, cool=int))
+
     def test_pickle(self):
         global Emp  # pickle wants to reference the class by name
         Emp = NamedTuple('Emp', [('name', str), ('id', int)])

src/typing.py

@@ -1801,31 +1801,66 @@ def new_user(user_class: Type[U]) -> U:
     """
 
 
-def NamedTuple(typename, fields):
-    """Typed version of namedtuple.
+def _make_nmtuple(name, types):
+    nm_tpl = collections.namedtuple(name, [n for n, t in types])
+    nm_tpl._field_types = dict(types)
+    try:
+        nm_tpl.__module__ = sys._getframe(2).f_globals.get('__name__', '__main__')
+    except (AttributeError, ValueError):
+        pass
+    return nm_tpl
 
-    Usage::
 
-        Employee = typing.NamedTuple('Employee', [('name', str), 'id', int)])
+if sys.version_info[:2] >= (3, 6):
+    class NamedTupleMeta(type):
 
-    This is equivalent to::
+        def __new__(cls, typename, bases, ns, *, _root=False):
+            if _root:
+                return super().__new__(cls, typename, bases, ns)
+            types = ns.get('__annotations__', {})
+            return _make_nmtuple(typename, types.items())
 
-        Employee = collections.namedtuple('Employee', ['name', 'id'])
+    class NamedTuple(metaclass=NamedTupleMeta, _root=True):
+        """Typed version of namedtuple.
 
-    The resulting class has one extra attribute: _field_types,
-    giving a dict mapping field names to types.  (The field names
-    are in the _fields attribute, which is part of the namedtuple
-    API.)
-    """
-    fields = [(n, t) for n, t in fields]
-    cls = collections.namedtuple(typename, [n for n, t in fields])
-    cls._field_types = dict(fields)
-    # Set the module to the caller's module (otherwise it'd be 'typing').
-    try:
-        cls.__module__ = sys._getframe(1).f_globals.get('__name__', '__main__')
-    except (AttributeError, ValueError):
-        pass
-    return cls
+        Usage::
+
+            class Employee(NamedTuple):
+                name: str
+                id: int
+
+        This is equivalent to::
+
+            Employee = collections.namedtuple('Employee', ['name', 'id'])
+
+        The resulting class has one extra attribute: _field_types,
+        giving a dict mapping field names to types.  (The field names
+        are in the _fields attribute, which is part of the namedtuple
+        API.) Backward-compatible usage::
+
+            Employee = NamedTuple('Employee', [('name', str), ('id', int)])
+        """
+
+        def __new__(self, typename, fields):
+            return _make_nmtuple(typename, fields)
+else:
+    def NamedTuple(typename, fields):
+        """Typed version of namedtuple.
+
+        Usage::
+
+            Employee = typing.NamedTuple('Employee', [('name', str), 'id', int)])
+
+        This is equivalent to::
+
+            Employee = collections.namedtuple('Employee', ['name', 'id'])
+
+        The resulting class has one extra attribute: _field_types,
+        giving a dict mapping field names to types.  (The field names
+        are in the _fields attribute, which is part of the namedtuple
+        API.)
+        """
+        return _make_nmtuple(typename, fields)
 
 
 def NewType(name, tp):