From 245a927382d90ff694dd8187611e855f33475a5b Mon Sep 17 00:00:00 2001
From: Yuyao Huang <huangyuyao@outlook.com>
Date: Fri, 19 Apr 2024 19:16:39 +0800
Subject: [PATCH] Compare()

---
 experiment.py                           |  3 ++-
 tests/test_conditions.py                | 19 +++++++++++++++++++
 trace_commentor/handlers/__init__.py    |  2 +-
 trace_commentor/handlers/expressions.py |  5 +++++
 4 files changed, 27 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_conditions.py

diff --git a/experiment.py b/experiment.py
index 7be429a..9d217eb 100644
--- a/experiment.py
+++ b/experiment.py
@@ -5,7 +5,8 @@ from tests.test_utils import *
     
 ])
 def function():
-    a, b = 3, 4
+    x = 2
+    print(x == 2)
 
 
 print(function())
diff --git a/tests/test_conditions.py b/tests/test_conditions.py
new file mode 100644
index 0000000..7d8a543
--- /dev/null
+++ b/tests/test_conditions.py
@@ -0,0 +1,19 @@
+from test_utils import *
+
+
+def test_constant():
+
+    @Commentor("<return>")
+    def target():
+        x = 2
+        print(x == 2)
+
+    asserteq_or_print(target(), '''
+    def target():
+        x = 2
+        print(x == 2)
+        """
+        True : x == 2
+        None : print(x == 2)
+        """
+''')
diff --git a/trace_commentor/handlers/__init__.py b/trace_commentor/handlers/__init__.py
index c2ddf19..fc28924 100644
--- a/trace_commentor/handlers/__init__.py
+++ b/trace_commentor/handlers/__init__.py
@@ -1,5 +1,5 @@
 from .definitions import FunctionDef, Return
 from .statements import Pass, Assign
-from .expressions import Expr, BinOp, Call
+from .expressions import Expr, BinOp, Call, Compare
 from .literals import Constant, Tuple
 from .variables import Name
diff --git a/trace_commentor/handlers/expressions.py b/trace_commentor/handlers/expressions.py
index 82e81ee..8113cc6 100644
--- a/trace_commentor/handlers/expressions.py
+++ b/trace_commentor/handlers/expressions.py
@@ -6,6 +6,11 @@ def BinOp(self, cmtor):
     cmtor.process(self.right)
     cmtor.append_comment(cmtor.eval(self))
 
+def Compare(self, cmtor):
+    for cmp in self.comparators:
+        cmtor.process(cmp)
+    cmtor.append_comment(cmtor.eval(self))
+
 def Call(self, cmtor):
     for arg in self.args:
         cmtor.process(arg)