From 5b5e9ce04f86a7cef4bcb91ac6bf7abc5ada174a Mon Sep 17 00:00:00 2001
From: pca006132 <jl@m-labs.hk>
Date: Mon, 4 Jan 2021 14:53:52 +0800
Subject: [PATCH] added magic methods

---
 nac3core/src/lib.rs           |  1 +
 nac3core/src/magic_methods.rs | 58 +++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 nac3core/src/magic_methods.rs

diff --git a/nac3core/src/lib.rs b/nac3core/src/lib.rs
index b3fcf1a5..323dd504 100644
--- a/nac3core/src/lib.rs
+++ b/nac3core/src/lib.rs
@@ -6,6 +6,7 @@ extern crate inkwell;
 extern crate rustpython_parser;
 
 pub mod inference_core;
+mod magic_methods;
 pub mod primitives;
 pub mod typedef;
 pub mod context;
diff --git a/nac3core/src/magic_methods.rs b/nac3core/src/magic_methods.rs
new file mode 100644
index 00000000..b0c248b4
--- /dev/null
+++ b/nac3core/src/magic_methods.rs
@@ -0,0 +1,58 @@
+use rustpython_parser::ast::{Comparison, Operator, UnaryOperator};
+
+pub fn binop_name(op: &Operator) -> &'static str {
+    match op {
+        Operator::Add => "__add__",
+        Operator::Sub => "__sub__",
+        Operator::Div => "__truediv__",
+        Operator::Mod => "__mod__",
+        Operator::Mult => "__mul__",
+        Operator::Pow => "__pow__",
+        Operator::BitOr => "__or__",
+        Operator::BitXor => "__xor__",
+        Operator::BitAnd => "__and__",
+        Operator::LShift => "__lshift__",
+        Operator::RShift => "__rshift__",
+        Operator::FloorDiv => "__floordiv__",
+        Operator::MatMult => "__matmul__",
+    }
+}
+
+pub fn binop_assign_name(op: &Operator) -> &'static str {
+    match op {
+        Operator::Add => "__iadd__",
+        Operator::Sub => "__isub__",
+        Operator::Div => "__itruediv__",
+        Operator::Mod => "__imod__",
+        Operator::Mult => "__imul__",
+        Operator::Pow => "__ipow__",
+        Operator::BitOr => "__ior__",
+        Operator::BitXor => "__ixor__",
+        Operator::BitAnd => "__iand__",
+        Operator::LShift => "__ilshift__",
+        Operator::RShift => "__irshift__",
+        Operator::FloorDiv => "__ifloordiv__",
+        Operator::MatMult => "__imatmul__",
+    }
+}
+
+pub fn unaryop_name(op: &UnaryOperator) -> &'static str {
+    match op {
+        UnaryOperator::Pos => "__pos__",
+        UnaryOperator::Neg => "__neg__",
+        UnaryOperator::Not => "__not__",
+        UnaryOperator::Inv => "__inv__",
+    }
+}
+
+pub fn comparison_name(op: &Comparison) -> Option<&'static str> {
+    match op {
+        Comparison::Less => Some("__lt__"),
+        Comparison::LessOrEqual => Some("__le__"),
+        Comparison::Greater => Some("__gt__"),
+        Comparison::GreaterOrEqual => Some("__ge__"),
+        Comparison::Equal => Some("__eq__"),
+        Comparison::NotEqual => Some("__ne__"),
+        _ => None,
+    }
+}