Skip to content

Commit

Permalink
Support typing.assert_type.
Browse files Browse the repository at this point in the history
* Aliases typing.assert_type to our assert_type pseudo-builtin.
* Removes single-argument assert_type form. This functionality didn't make it
  into the typing library variant. I don't think it's used much anyway,
  given that pytype_extensions.assert_type doesn't support it at runtime, and
  no one's ever noticed.

PiperOrigin-RevId: 577000283
  • Loading branch information
rchen152 committed Oct 27, 2023
1 parent c452719 commit 5b5a22d
Show file tree
Hide file tree
Showing 6 changed files with 15 additions and 26 deletions.
8 changes: 1 addition & 7 deletions pytype/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -1250,16 +1250,10 @@ def reveal_type(self, stack, node, var):
self.error(stack, self._var_to_printed_type(var, node))

@_error_name("assert-type")
def assert_type(self, stack, node, var, typ=None):
def assert_type(self, stack, node, var, typ):
"""Check that a variable type matches its expected value."""
actual = self._var_to_printed_type(var, node)

# assert_type(x) checks that x is not Any
if typ is None:
if actual in ("Any", "typing.Any"):
self.error(stack, f"Asserted type was {actual}")
return

try:
expected = abstract_utils.get_atomic_python_constant(typ, str)
except abstract_utils.ConversionError:
Expand Down
1 change: 1 addition & 0 deletions pytype/overlays/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,7 @@ py_library(
.named_tuple
.overlay
.overlay_utils
.special_builtins
.typed_dict
pytype.utils
pytype.abstract.abstract
Expand Down
5 changes: 1 addition & 4 deletions pytype/overlays/special_builtins.py
Original file line number Diff line number Diff line change
Expand Up @@ -553,10 +553,7 @@ class AssertType(BuiltinFunction):
_NAME = "assert_type"

def call(self, node, func, args, alias_map=None):
if len(args.posargs) == 1:
a, = args.posargs
t = None
elif len(args.posargs) == 2:
if len(args.posargs) == 2:
a, t = args.posargs
else:
raise function.WrongArgCount(self._SIGNATURE, args, self.ctx)
Expand Down
5 changes: 4 additions & 1 deletion pytype/overlays/typing_overlay.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from pytype.overlays import named_tuple
from pytype.overlays import overlay
from pytype.overlays import overlay_utils
from pytype.overlays import special_builtins
from pytype.overlays import typed_dict
from pytype.pytd import pep484
from pytype.pytd import pytd
Expand Down Expand Up @@ -627,7 +628,6 @@ def build_re_member(ctx, module):
"TypeVarTuple": (3, 11),
"Unpack": (3, 11),
"assert_never": (3, 11),
"assert_type": (3, 11),
"reveal_type": (3, 11),
}

Expand Down Expand Up @@ -658,6 +658,9 @@ def build_re_member(ctx, module):
"TypedDict": (overlay.drop_module(typed_dict.TypedDictBuilder), (3, 8)),
"Union": (overlay.drop_module(Union), None),
"TYPE_CHECKING": (overlay.drop_module(build_typechecking), None),
"assert_type": (
overlay.add_name("assert_type", special_builtins.AssertType.make_alias),
(3, 11)),
"cast": (overlay.add_name("cast", Cast.make), None),
"clear_overloads": (_builder_from_name("clear_overloads"), (3, 11)),
"dataclass_transform": (overlay.add_name(
Expand Down
18 changes: 6 additions & 12 deletions pytype/tests/test_errors2.py
Original file line number Diff line number Diff line change
Expand Up @@ -345,31 +345,27 @@ def test_assert_type(self):
class A: pass
def f(x: int, y: str, z):
assert_type(x, int)
assert_type(y, int) # assert-type[e1]
assert_type(z) # assert-type[e2]
assert_type(y, int) # assert-type[e]
if __random__:
x = A()
assert_type(x, Union[A, int])
""")
self.assertErrorSequences(errors, {
"e1": ["Expected", "int", "Actual", "str"],
"e2": ["type was Any"]
"e": ["Expected", "int", "Actual", "str"],
})

def test_assert_type_str(self):
_, errors = self.InferWithErrors("""
class A: pass
def f(x: int, y: str, z):
assert_type(x, 'int')
assert_type(y, 'int') # assert-type[e1]
assert_type(z) # assert-type[e2]
assert_type(y, 'int') # assert-type[e]
if __random__:
x = A()
assert_type(x, 'Union[A, int]')
""")
self.assertErrorSequences(errors, {
"e1": ["Expected", "int", "Actual", "str"],
"e2": ["type was Any"]
"e": ["Expected", "int", "Actual", "str"],
})

def test_assert_type_import(self):
Expand All @@ -383,15 +379,13 @@ def assert_type(*args): ...
class A: pass
def f(x: int, y: str, z):
assert_type(x, int)
assert_type(y, int) # assert-type[e1]
assert_type(z) # assert-type[e2]
assert_type(y, int) # assert-type[e]
if __random__:
x = A()
assert_type(x, Union[A, int])
""", pythonpath=[d.path])
self.assertErrorSequences(errors, {
"e1": ["Expected", "int", "Actual", "str"],
"e2": ["type was Any"]
"e": ["Expected", "int", "Actual", "str"],
})

def test_combine_containers(self):
Expand Down
4 changes: 2 additions & 2 deletions pytype/tests/test_pattern_matching.py
Original file line number Diff line number Diff line change
Expand Up @@ -718,8 +718,8 @@ def f(x: A | B, y: A | B):
assert_type(c, int)
assert_type(d, str)
case _:
assert_type(x)
assert_type(y)
assert_type(x, B)
assert_type(y, A)
""")

def test_builtin(self):
Expand Down

0 comments on commit 5b5a22d

Please sign in to comment.