mypy_extensions.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. """Defines experimental extensions to the standard "typing" module that are
  2. supported by the mypy typechecker.
  3. Example usage:
  4. from mypy_extensions import TypedDict
  5. """
  6. from typing import Any
  7. # NOTE: This module must support Python 2.7 in addition to Python 3.x
  8. import sys
  9. # _type_check is NOT a part of public typing API, it is used here only to mimic
  10. # the (convenient) behavior of types provided by typing module.
  11. from typing import _type_check # type: ignore
  12. def _check_fails(cls, other):
  13. try:
  14. if sys._getframe(1).f_globals['__name__'] not in ['abc', 'functools', 'typing']:
  15. # Typed dicts are only for static structural subtyping.
  16. raise TypeError('TypedDict does not support instance and class checks')
  17. except (AttributeError, ValueError):
  18. pass
  19. return False
  20. def _dict_new(cls, *args, **kwargs):
  21. return dict(*args, **kwargs)
  22. def _typeddict_new(cls, _typename, _fields=None, **kwargs):
  23. total = kwargs.pop('total', True)
  24. if _fields is None:
  25. _fields = kwargs
  26. elif kwargs:
  27. raise TypeError("TypedDict takes either a dict or keyword arguments,"
  28. " but not both")
  29. return _TypedDictMeta(_typename, (), {'__annotations__': dict(_fields),
  30. '__total__': total})
  31. class _TypedDictMeta(type):
  32. def __new__(cls, name, bases, ns, total=True):
  33. # Create new typed dict class object.
  34. # This method is called directly when TypedDict is subclassed,
  35. # or via _typeddict_new when TypedDict is instantiated. This way
  36. # TypedDict supports all three syntaxes described in its docstring.
  37. # Subclasses and instances of TypedDict return actual dictionaries
  38. # via _dict_new.
  39. ns['__new__'] = _typeddict_new if name == 'TypedDict' else _dict_new
  40. tp_dict = super(_TypedDictMeta, cls).__new__(cls, name, (dict,), ns)
  41. try:
  42. # Setting correct module is necessary to make typed dict classes pickleable.
  43. tp_dict.__module__ = sys._getframe(2).f_globals.get('__name__', '__main__')
  44. except (AttributeError, ValueError):
  45. pass
  46. anns = ns.get('__annotations__', {})
  47. msg = "TypedDict('Name', {f0: t0, f1: t1, ...}); each t must be a type"
  48. anns = {n: _type_check(tp, msg) for n, tp in anns.items()}
  49. for base in bases:
  50. anns.update(base.__dict__.get('__annotations__', {}))
  51. tp_dict.__annotations__ = anns
  52. if not hasattr(tp_dict, '__total__'):
  53. tp_dict.__total__ = total
  54. return tp_dict
  55. __instancecheck__ = __subclasscheck__ = _check_fails
  56. TypedDict = _TypedDictMeta('TypedDict', (dict,), {})
  57. TypedDict.__module__ = __name__
  58. TypedDict.__doc__ = \
  59. """A simple typed name space. At runtime it is equivalent to a plain dict.
  60. TypedDict creates a dictionary type that expects all of its
  61. instances to have a certain set of keys, with each key
  62. associated with a value of a consistent type. This expectation
  63. is not checked at runtime but is only enforced by typecheckers.
  64. Usage::
  65. Point2D = TypedDict('Point2D', {'x': int, 'y': int, 'label': str})
  66. a: Point2D = {'x': 1, 'y': 2, 'label': 'good'} # OK
  67. b: Point2D = {'z': 3, 'label': 'bad'} # Fails type check
  68. assert Point2D(x=1, y=2, label='first') == dict(x=1, y=2, label='first')
  69. The type info could be accessed via Point2D.__annotations__. TypedDict
  70. supports two additional equivalent forms::
  71. Point2D = TypedDict('Point2D', x=int, y=int, label=str)
  72. class Point2D(TypedDict):
  73. x: int
  74. y: int
  75. label: str
  76. The latter syntax is only supported in Python 3.6+, while two other
  77. syntax forms work for Python 2.7 and 3.2+
  78. """
  79. # Argument constructors for making more-detailed Callables. These all just
  80. # return their type argument, to make them complete noops in terms of the
  81. # `typing` module.
  82. def Arg(type=Any, name=None):
  83. """A normal positional argument"""
  84. return type
  85. def DefaultArg(type=Any, name=None):
  86. """A positional argument with a default value"""
  87. return type
  88. def NamedArg(type=Any, name=None):
  89. """A keyword-only argument"""
  90. return type
  91. def DefaultNamedArg(type=Any, name=None):
  92. """A keyword-only argument with a default value"""
  93. return type
  94. def VarArg(type=Any):
  95. """A *args-style variadic positional argument"""
  96. return type
  97. def KwArg(type=Any):
  98. """A **kwargs-style variadic keyword argument"""
  99. return type
  100. # Return type that indicates a function does not return
  101. class NoReturn: pass
  102. def trait(cls):
  103. return cls
  104. # TODO: We may want to try to properly apply this to any type
  105. # variables left over...
  106. class _FlexibleAliasClsApplied:
  107. def __init__(self, val):
  108. self.val = val
  109. def __getitem__(self, args):
  110. return self.val
  111. class _FlexibleAliasCls:
  112. def __getitem__(self, args):
  113. return _FlexibleAliasClsApplied(args[-1])
  114. FlexibleAlias = _FlexibleAliasCls()