diff --git a/mypy/semanal.py b/mypy/semanal.py index 8d9b001ae775..803813bd40eb 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -3406,11 +3406,19 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None: self.process__slots__(s) def is_sentinel_declaration(self, s: AssignmentStmt) -> bool: - """Does this assignment define a PEP 661 sentinel singleton?""" + """Does this assignment define a PEP 661 sentinel singleton? + + This includes both the original `NAME = Sentinel("NAME")` call and a + plain alias of an existing sentinel, e.g. `ALIAS = NAME` or + `ALIAS = mod.NAME`, so that re-exports of a sentinel keep working as + the same sentinel type. + """ if self.is_nested_within_func_scope() or s.unanalyzed_type is not None: return False if len(s.lvalues) != 1 or not isinstance(s.lvalues[0], NameExpr): return False + if isinstance(s.rvalue, RefExpr): + return isinstance(s.rvalue.node, Var) and s.rvalue.node.is_sentinel if not isinstance(s.rvalue, CallExpr): return False call = s.rvalue @@ -3430,7 +3438,13 @@ def setup_sentinel_var(self, s: AssignmentStmt) -> None: lvalue.is_special_form = True var = lvalue.node var.is_sentinel = True - typ = self.sentinel_type_for_var(var, s.rvalue) + if isinstance(s.rvalue, RefExpr): + # An alias of an existing sentinel: reuse its already-computed type. + assert isinstance(s.rvalue.node, Var) + typ = get_proper_type(s.rvalue.node.type) + assert isinstance(typ, LiteralType) + else: + typ = self.sentinel_type_for_var(var, s.rvalue) if typ is not None: s.type = typ diff --git a/test-data/unit/check-sentinels.test b/test-data/unit/check-sentinels.test index 7a85cda95e81..8b3aa79fca95 100644 --- a/test-data/unit/check-sentinels.test +++ b/test-data/unit/check-sentinels.test @@ -176,6 +176,22 @@ def func(x: int | BUILTIN | TYPEXT) -> None: [builtins fixtures/sentinel.pyi] [case testSentinelReassignmentIsNotTypeAlias] +from typing_extensions import sentinel + +MISSING = sentinel("MISSING") +ALIAS = MISSING + +# Not a type alias (i.e. not a special form) +reveal_type(ALIAS) # N: Revealed type is "MISSING" + +def func(x: int | MISSING = MISSING) -> None: + pass + +func(MISSING) +func(ALIAS) +[builtins fixtures/tuple.pyi] + +[case testSentinelReassignmentIsSameSentinel] from typing_extensions import sentinel, assert_type MISSING = sentinel("MISSING") @@ -184,16 +200,38 @@ ALIAS = MISSING # The value still identifies as the same sentinel... assert_type(ALIAS, MISSING) -def func(x: int | MISSING = MISSING) -> None: +def func(x: ALIAS | int = ALIAS) -> None: pass func(MISSING) func(ALIAS) -# ...but the reassignment does not make ALIAS usable as a type alias. -def uses_alias_as_type(x: ALIAS) -> None: # E: Variable "__main__.ALIAS" is not valid as a type \ - # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases +# ...and the reassignment keeps ALIAS usable as a type expression too. +def uses_alias_as_type(x: ALIAS) -> None: + pass + +uses_alias_as_type(MISSING) +uses_alias_as_type(ALIAS) +[builtins fixtures/tuple.pyi] + +[case testSentinelReassignmentAcrossModules] +from other import MISSING + +def func(x: int | MISSING = MISSING) -> None: pass + +func(MISSING) + +[file other.py] +from typing_extensions import sentinel +import lib + +MISSING = lib.MISSING + +[file lib.py] +from typing_extensions import sentinel + +MISSING = sentinel("MISSING") [builtins fixtures/tuple.pyi] [case testSentinelPreservedThroughGenericSubstitution]