Also support set expressions.

Author: tyralla

docs/source/literal_types.rst

@@ -368,8 +368,8 @@ without a value:
       elif x == 'two':
           return False
 
-For the sake of brevity, you can use the ``in`` operator in combination with list or tuple
-expressions (lists or tuples created "on the fly"):
+For the sake of brevity, you can use the ``in`` operator in combination with
+list, set, or tuple expressions (lists, sets, or tuples created "on the fly"):
 
 .. code-block:: python
 

mypy/checker.py

@@ -117,6 +117,7 @@
     RaiseStmt,
     RefExpr,
     ReturnStmt,
+    SetExpr,
     StarExpr,
     Statement,
     StrExpr,
@@ -4685,11 +4686,11 @@ def check_return_stmt(self, s: ReturnStmt) -> None:
                 if self.in_checked_function():
                     self.fail(message_registry.RETURN_VALUE_EXPECTED, s)
 
-    def _make_tupleexpr_with_literals_narrowable_by_using_in(self, e: Expression) -> Expression:
+    def _transform_sequence_expressions_for_narrowing_with_in(self, e: Expression) -> Expression:
         """
         Transform an expression like
 
-        (x is None) and (x in (1, 2)) and (x not in (3, 4))
+        (x is None) and (x in (1, 2)) and (x not in [3, 4])
 
         into
 
@@ -4700,15 +4701,15 @@ def _make_tupleexpr_with_literals_narrowable_by_using_in(self, e: Expression) ->
         implement additional narrowing logic.
         """
         if isinstance(e, OpExpr):
-            e.left = self._make_tupleexpr_with_literals_narrowable_by_using_in(e.left)
-            e.right = self._make_tupleexpr_with_literals_narrowable_by_using_in(e.right)
+            e.left = self._transform_sequence_expressions_for_narrowing_with_in(e.left)
+            e.right = self._transform_sequence_expressions_for_narrowing_with_in(e.right)
             return e
 
         if not (
             isinstance(e, ComparisonExpr)
             and isinstance(left := e.operands[0], NameExpr)
             and ((op_in := e.operators[0]) in ("in", "not in"))
-            and isinstance(litu := e.operands[1], (ListExpr, TupleExpr))
+            and isinstance(litu := e.operands[1], (ListExpr, SetExpr, TupleExpr))
         ):
             return e
 
@@ -4748,7 +4749,7 @@ def visit_if_stmt(self, s: IfStmt) -> None:
         with self.binder.frame_context(can_skip=False, conditional_frame=True, fall_through=0):
             for e, b in zip(s.expr, s.body):
 
-                e = self._make_tupleexpr_with_literals_narrowable_by_using_in(e)
+                e = self._transform_sequence_expressions_for_narrowing_with_in(e)
 
                 t = get_proper_type(self.expr_checker.accept(e))
 

test-data/unit/check-narrowing.test

@@ -2334,7 +2334,7 @@ def f(x: C) -> None:
 f(C(5))
 [builtins fixtures/primitives.pyi]
 
-[case testNarrowLiteralsInListOrTupleExpression]
+[case testNarrowLiteralsInListOrSetOrTupleExpression]
 # flags: --warn-unreachable
 
 from typing import Optional
@@ -2347,7 +2347,7 @@ def f(v: Optional[Literal[1, 2, 3, 4]]) -> None:
         reveal_type(v)  # N: Revealed type is "Union[Literal[1], Literal[2]]"
     elif v in [1]:
         reveal_type(v)  # E: Statement is unreachable
-    elif v is None or v in [3, x]:
+    elif v is None or v in {3, x}:
         reveal_type(v)  # N: Revealed type is "Union[Literal[3], Literal[4], None]"
     elif v in ():
         reveal_type(v)  # E: Statement is unreachable
@@ -2356,7 +2356,7 @@ def f(v: Optional[Literal[1, 2, 3, 4]]) -> None:
     reveal_type(v)  # N: Revealed type is "Union[Literal[1], Literal[2], Literal[3], Literal[4], None]"
 [builtins fixtures/primitives.pyi]
 
-[case testNarrowLiteralsNotInListOrTupleExpression]
+[case testNarrowLiteralsNotInListOrSetOrTupleExpression]
 # flags: --warn-unreachable
 
 from typing import Optional
@@ -2365,7 +2365,7 @@ from typing_extensions import Literal
 x: int
 
 def f(v: Optional[Literal[1, 2, 3, 4, 5]]) -> None:
-    if v not in (0, 1, 2, 3):
+    if v not in {0, 1, 2, 3}:
         reveal_type(v)  # N: Revealed type is "Union[Literal[4], Literal[5], None]"
     elif v not in [1, 2, 3, 4]:  # E: Right operand of "and" is never evaluated
         reveal_type(v)  # E: Statement is unreachable
@@ -2378,7 +2378,7 @@ def f(v: Optional[Literal[1, 2, 3, 4, 5]]) -> None:
     reveal_type(v)  # N: Revealed type is "Union[Literal[1], Literal[2], Literal[3], Literal[4], Literal[5], None]"
 [builtins fixtures/primitives.pyi]
 
-[case testNarrowEnumsInListOrTupleExpression]
+[case testNarrowEnumsInListOrSetOrTupleExpression]
 from enum import Enum
 from typing import Final
 
@@ -2397,7 +2397,7 @@ def f(v: E) -> None:
         reveal_type(v)  # N: Revealed type is "Union[Literal[__main__.E.A], Literal[__main__.E.B]]"
     elif v in [E.A]:
         reveal_type(v)
-    elif v in (C,):
+    elif v in {C}:
         reveal_type(v)  # N: Revealed type is "Literal[__main__.E.C]"
     elif v in ():
         reveal_type(v)
@@ -2406,7 +2406,7 @@ def f(v: E) -> None:
     reveal_type(v)  # N: Revealed type is "__main__.E"
 [builtins fixtures/primitives.pyi]
 
-[case testNarrowEnumsNotInListOrTupleExpression]
+[case testNarrowEnumsNotInListOrSetOrTupleExpression]
 from enum import Enum
 from typing import Final
 
@@ -2426,7 +2426,7 @@ def f(v: E) -> None:
         reveal_type(v)  # N: Revealed type is "Union[Literal[__main__.E.D], Literal[__main__.E.E]]"
     elif v not in [E.A, E.B, E.C, E.C]:
         reveal_type(v)
-    elif v not in (C,):
+    elif v not in {C}:
         reveal_type(v)  # N: Revealed type is "Union[Literal[__main__.E.A], Literal[__main__.E.B]]"
     elif v not in []:
         reveal_type(v)  # N: Revealed type is "Literal[__main__.E.C]"