Skip to content

Commit edd5b87

Browse files
committed
Add imag function for complex
1 parent fef9fbd commit edd5b87

File tree

2 files changed

+27
-0
lines changed

2 files changed

+27
-0
lines changed

src/lpython/semantics/python_comptime_eval.h

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ struct PythonIntrinsicProcedures {
4444
{"hex", {m_builtin, &eval_hex}},
4545
{"oct", {m_builtin, &eval_oct}},
4646
{"complex", {m_builtin, &eval_complex}},
47+
{"imag", {m_builtin, &eval_imag}},
4748
{"divmod", {m_builtin, &eval_divmod}},
4849
};
4950
}
@@ -494,6 +495,24 @@ struct PythonIntrinsicProcedures {
494495
return ASR::down_cast<ASR::expr_t>(make_ConstantComplex_t(al, loc, c1, c2, type));
495496
}
496497

498+
static ASR::expr_t *eval_imag(Allocator &al, const Location &loc,
499+
Vec<ASR::expr_t*> &args
500+
) {
501+
LFORTRAN_ASSERT(ASRUtils::all_args_evaluated(args));
502+
if (args.size() != 1) {
503+
throw SemanticError("Intrinsic imag function accepts exactly 1 argument", loc);
504+
}
505+
ASR::expr_t* imag_arg = args[0];
506+
ASR::ttype_t* t = ASRUtils::expr_type(args[0]);
507+
if (LFortran::ASR::is_a<LFortran::ASR::Complex_t>(*t)) {
508+
double im = ASR::down_cast<ASR::ConstantComplex_t>(imag_arg)->m_im;
509+
double result = im;
510+
return ASR::down_cast<ASR::expr_t>(ASR::make_ConstantReal_t(al, loc, result, t));
511+
} else {
512+
throw SemanticError("Argument of the aimag() function must be Complex", loc);
513+
}
514+
}
515+
497516
static ASR::expr_t *eval_divmod(Allocator &al, const Location &loc, Vec<ASR::expr_t *> &args) {
498517
LFORTRAN_ASSERT(ASRUtils::all_args_evaluated(args));
499518
if (args.size() != 2) {

src/runtime/lpython_builtin.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,3 +288,11 @@ def lbound(x: i32[:], dim: i32) -> i32:
288288

289289
def ubound(x: i32[:], dim: i32) -> i32:
290290
pass
291+
292+
293+
@overload
294+
def imag(x: c64) -> f64:
295+
pass
296+
297+
@overload
298+
def imag(x: c32) -> f32:

0 commit comments

Comments
 (0)