ASR: Support array.size attribute

virendrakabra14 · May 11, 2023 · d1535f4 · d1535f4
1 parent 9b8fa1e
commit d1535f4
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp
@@ -5137,7 +5137,13 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
 
     void visit_AttributeUtil(ASR::ttype_t* type, char* attr_char,
                              ASR::symbol_t *t, const Location& loc) {
-        if (ASRUtils::is_complex(*type)) {
+        if (ASRUtils::is_array(type)) {
+            std::string attr = attr_char;
+            ASR::expr_t *se = ASR::down_cast<ASR::expr_t>(ASR::make_Var_t(al, loc, t));
+            Vec<ASR::expr_t*> args;
+            args.reserve(al, 0);
+            handle_attribute(se, attr, loc, args);
+        } else if (ASRUtils::is_complex(*type)) {
             std::string attr = attr_char;
             if (attr == "imag") {
                 ASR::expr_t *val = ASR::down_cast<ASR::expr_t>(ASR::make_Var_t(al, loc, t));
@@ -6705,7 +6711,7 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                     ASRUtils::IntrinsicFunctionRegistry::get_create_function(call_name);
                 Vec<ASR::expr_t*> args_; args_.reserve(al, x.n_args);
                 visit_expr_list(x.m_args, x.n_args, args_);
-                if (ASRUtils::is_array(ASRUtils::expr_type(args_[0])) && 
+                if (ASRUtils::is_array(ASRUtils::expr_type(args_[0])) &&
                     imported_functions[call_name] == "math" ) {
                     throw SemanticError("Function '" + call_name + "' does not accept vector values",
                         x.base.base.loc);

diff --git a/src/lpython/semantics/python_attribute_eval.h b/src/lpython/semantics/python_attribute_eval.h
@@ -20,6 +20,7 @@ struct AttributeHandler {
     AttributeHandler() {
         attribute_map = {
             {"int@bit_length", &eval_int_bit_length},
+            {"array@size", &eval_array_size},
             {"list@append", &eval_list_append},
             {"list@remove", &eval_list_remove},
             {"list@count", &eval_list_count},
@@ -42,6 +43,8 @@ struct AttributeHandler {
             return "set";
         } else if (ASR::is_a<ASR::Dict_t>(*t)) {
             return "dict";
+        } else if (ASRUtils::is_array(t)) {
+            return "array";
         } else if (ASR::is_a<ASR::Integer_t>(*t)) {
             return "int";
         }
@@ -55,7 +58,7 @@ struct AttributeHandler {
         if (class_name == "") {
             throw SemanticError("Type name is not implemented yet.", loc);
         }
-        std::string key = get_type_name(type) + "@" + attr_name;
+        std::string key = class_name + "@" + attr_name;
         auto search = attribute_map.find(key);
         if (search != attribute_map.end()) {
             attribute_eval_callback cb = search->second;
@@ -77,6 +80,15 @@ struct AttributeHandler {
         return ASR::make_IntegerBitLen_t(al, loc, s, int_type, nullptr);
     }
 
+    static ASR::asr_t* eval_array_size(ASR::expr_t *s, Allocator &al, const Location &loc,
+            Vec<ASR::expr_t*> &args, diag::Diagnostics &/*diag*/) {
+        if (args.size() != 0) {
+            throw SemanticError("array.size() takes no arguments", loc);
+        }
+        ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, loc, 4, nullptr, 0));
+        return ASR::make_ArraySize_t(al, loc, s, nullptr, int_type, nullptr);
+    }
+
     static ASR::asr_t* eval_list_append(ASR::expr_t *s, Allocator &al, const Location &loc,
             Vec<ASR::expr_t*> &args, diag::Diagnostics &diag) {
         if (args.size() != 1) {