From 92c92f6a4c7d4f2c890176a793b1c840d8f0e68e Mon Sep 17 00:00:00 2001
From: pca006132 <jl@m-labs.hk>
Date: Fri, 18 Dec 2020 16:54:47 +0800
Subject: [PATCH] added operator override test

---
 toy-impl/test_expr.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/toy-impl/test_expr.py b/toy-impl/test_expr.py
index bd97f3f..e12e006 100644
--- a/toy-impl/test_expr.py
+++ b/toy-impl/test_expr.py
@@ -3,6 +3,7 @@ from type_def import *
 from inference import *
 from helper import *
 from parse_expr import *
+from top_level import *
 
 types = {
     'int32': PrimitiveType('int32'),
@@ -71,3 +72,31 @@ test_expr('[[1], [2], [3]][a]', {'a': i32})
 test_expr('a == a == a', {'a': I})
 test_expr('a == a and 1 == 2', {'a': I})
 
+test_classes = """
+class Foo:
+    a: int32
+    def __init__(self, a: int32):
+        self.a = a
+
+    def __add__(self, other: Foo) -> Foo:
+        return Foo(self.a + other.a)
+
+    def __sub__(self, other: Foo) -> Foo:
+        return Foo(self.a - other.a)
+
+    def __lt__(self, other: Foo) -> bool:
+        pass
+
+    def __le__(self, other: Foo) -> bool:
+        pass
+
+    def __gt__(self, other: Foo) -> bool:
+        pass
+
+    def __ge__(self, other: Foo) -> bool:
+        pass
+"""
+ctx, _ = parse_top_level(ctx, ast.parse(test_classes))
+test_expr('Foo(1) + Foo(1)', {})
+test_expr('Foo(1) + Foo(1) < Foo(2) + Foo(3) < Foo(4)', {})
+