Add imag function for complex

Smit-create · Smit-create · commit 0196153c8fb4 · 2022-03-24T10:36:24.000+05:30
diff --git a/src/lpython/semantics/python_comptime_eval.h b/src/lpython/semantics/python_comptime_eval.h
@@ -44,6 +44,7 @@ struct PythonIntrinsicProcedures {
             {"hex", {m_builtin, &eval_hex}},
             {"oct", {m_builtin, &eval_oct}},
             {"complex", {m_builtin, &eval_complex}},
+            {"imag", {m_builtin, &eval_imag}},
             {"divmod", {m_builtin, &eval_divmod}},
         };
     }
@@ -503,6 +504,24 @@ struct PythonIntrinsicProcedures {
         return ASR::down_cast<ASR::expr_t>(make_ConstantComplex_t(al, loc, c1, c2, type));
     }
 
+    static ASR::expr_t *eval_imag(Allocator &al, const Location &loc,
+            Vec<ASR::expr_t*> &args
+            ) {
+        LFORTRAN_ASSERT(ASRUtils::all_args_evaluated(args));
+        if (args.size() != 1) {
+            throw SemanticError("Intrinsic imag function accepts exactly 1 argument", loc);
+        }
+        ASR::expr_t* imag_arg = args[0];
+        ASR::ttype_t* t = ASRUtils::expr_type(args[0]);
+        if (LFortran::ASR::is_a<LFortran::ASR::Complex_t>(*t)) {
+            double im = ASR::down_cast<ASR::ConstantComplex_t>(imag_arg)->m_im;
+            double result = im;
+            return ASR::down_cast<ASR::expr_t>(ASR::make_ConstantReal_t(al, loc, result, t));
+        } else {
+            throw SemanticError("Argument of the aimag() function must be Complex", loc);
+        }
+    }
+
     static ASR::expr_t *eval_divmod(Allocator &al, const Location &loc, Vec<ASR::expr_t *> &args) {
         LFORTRAN_ASSERT(ASRUtils::all_args_evaluated(args));
         if (args.size() != 2) {
diff --git a/src/runtime/lpython_builtin.py b/src/runtime/lpython_builtin.py
@@ -288,3 +288,11 @@ def lbound(x: i32[:], dim: i32) -> i32:
 
 def ubound(x: i32[:], dim: i32) -> i32:
     pass
+
+
+@overload
+def imag(x: c64) -> f64:
+    pass
+
+@overload
+def imag(x: c32) -> f32:
