Support for "in" comparison operator

scnerd · Jan 2, 2021 · 9fe14bc · 9fe14bc
1 parent dd4a549
commit 9fe14bc
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git 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
@@ -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)