diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt index d6ca666674..513b825c29 100644 --- a/integration_tests/CMakeLists.txt +++ b/integration_tests/CMakeLists.txt @@ -424,6 +424,7 @@ RUN(NAME array_01 LABELS cpython llvm wasm c) RUN(NAME array_02 LABELS cpython wasm c) RUN(NAME array_03 LABELS cpython llvm c) RUN(NAME array_04 LABELS cpython llvm c) +RUN(NAME array_05 LABELS cpython llvm c) RUN(NAME bindc_01 LABELS cpython llvm c) RUN(NAME bindc_02 LABELS cpython llvm c) RUN(NAME bindc_04 LABELS llvm c NOFAST) diff --git a/integration_tests/array_05.py b/integration_tests/array_05.py new file mode 100644 index 0000000000..4c18576cac --- /dev/null +++ b/integration_tests/array_05.py @@ -0,0 +1,39 @@ +from lpython import i32, f64, Array +from numpy import empty, int32, float64 + + +def test_1(): + y: Array[f64, 3] = empty([3], dtype=float64) + y[0] = 3.14 + y[1] = -4.14 + y[2] = 100.100 + + print(y) + assert abs(y[0] - (3.14)) <= 1e-6 + assert abs(y[1] - (-4.14)) <= 1e-6 + assert abs(y[2] - (100.100)) <= 1e-6 + +def test_2(): + x: Array[i32, 2, 3] = empty([2, 3], dtype=int32) + + x[0, 0] = 5 + x[0, 1] = -10 + x[0, 2] = 15 + x[1, 0] = 4 + x[1, 1] = -14 + x[1, 2] = 100 + + print(x) + assert x[0, 0] == 5 + assert x[0, 1] == -10 + assert x[0, 2] == 15 + assert x[1, 0] == 4 + assert x[1, 1] == -14 + assert x[1, 2] == 100 + + +def main0(): + test_1() + test_2() + +main0() diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp index c7376cb0dc..95325f4048 100644 --- a/src/lpython/semantics/python_ast_to_asr.cpp +++ b/src/lpython/semantics/python_ast_to_asr.cpp @@ -1851,17 +1851,35 @@ class CommonVisitor : public AST::BaseVisitor { is_allocatable, raise_error, abi, is_argument); return ASRUtils::TYPE(ASR::make_Const_t(al, loc, type)); } else { + AST::expr_t* dim_info = s->m_slice; + + if (var_annotation == "Array") { + LCOMPILERS_ASSERT(AST::is_a(*s->m_slice)); + AST::Tuple_t *t = AST::down_cast(s->m_slice); + LCOMPILERS_ASSERT(t->n_elts >= 2); + LCOMPILERS_ASSERT(AST::is_a(*t->m_elts[0])); + var_annotation = AST::down_cast(t->m_elts[0])->m_id; + Vec dims; + dims.reserve(al, 0); + for (size_t i = 1; i < t->n_elts; i++) { + dims.push_back(al, t->m_elts[i]); + } + AST::ast_t* dim_tuple = AST::make_Tuple_t(al, t->base.base.loc, dims.p, dims.size(), + AST::expr_contextType::Load); + dim_info = AST::down_cast(dim_tuple); + } + ASR::ttype_t* type = get_type_from_var_annotation(var_annotation, annotation.base.loc, dims, m_args, n_args, raise_error, abi, is_argument); - if (AST::is_a(*s->m_slice)) { + if (AST::is_a(*dim_info)) { ASR::dimension_t dim; dim.loc = loc; dim.m_start = nullptr; dim.m_length = nullptr; dims.push_back(al, dim); - } else if( is_runtime_array(s->m_slice) ) { - AST::Tuple_t* tuple_multidim = AST::down_cast(s->m_slice); + } else if( is_runtime_array(dim_info) ) { + AST::Tuple_t* tuple_multidim = AST::down_cast(dim_info); for( size_t i = 0; i < tuple_multidim->n_elts; i++ ) { if( AST::is_a(*tuple_multidim->m_elts[i]) ) { ASR::dimension_t dim; @@ -1872,7 +1890,7 @@ class CommonVisitor : public AST::BaseVisitor { } } } else { - this->visit_expr(*s->m_slice); + this->visit_expr(*dim_info); ASR::expr_t *value = ASRUtils::EXPR(tmp); fill_dims_for_asr_type(dims, value, loc); } diff --git a/src/runtime/lpython/lpython.py b/src/runtime/lpython/lpython.py index 073c8a4d3f..fa07f1e143 100644 --- a/src/runtime/lpython/lpython.py +++ b/src/runtime/lpython/lpython.py @@ -77,6 +77,9 @@ def __init__(self, type, dims): self._type = type self._dims = dims + def __class_getitem__(self, params): + return Array(params[0], params[1:]) + i1 = Type("i1") i8 = Type("i8") i16 = Type("i16")