diff --git a/pragma/core/resolve/__init__.py b/pragma/core/resolve/__init__.py index b21aa81..36681f0 100644 --- a/pragma/core/resolve/__init__.py +++ b/pragma/core/resolve/__init__.py @@ -75,6 +75,8 @@ def resolve_name_or_attribute(node, ctxt): ast.LtE: lambda a, b: a <= b, ast.Gt: lambda a, b: a > b, ast.GtE: lambda a, b: a >= b, + ast.In: lambda a, b: a in b, + ast.NotIn: lambda a, b: a not in b, } try: diff --git a/tests/test_collapse_literals.py b/tests/test_collapse_literals.py index 0c9deb9..8d45075 100644 --- a/tests/test_collapse_literals.py +++ b/tests/test_collapse_literals.py @@ -637,3 +637,29 @@ def f(x): yield 0 ''' self.assertSourceEqual(f, result) + + def test_collapse_InOp(self): + lst = ['a', 'b', object()] + dct = dict(a=1, b=2) + + @pragma.collapse_literals() + def f(): + if 'a' in lst: + yield 0 + if 'v' in lst: + unreachable + if 'b' not in lst: + unreachable + yield dct['a'] + if 'b' in dct: + yield 2 + # if 2 in dct.values(): # TODO: support this. Problem is that values is not a pure function. + # yield 2 + + result = ''' + def f(): + yield 0 + yield 1 + yield 2 + ''' + self.assertSourceEqual(f, result)