Add support for dict pop in python_attribute_eval.h

namannimmo10 · namannimmo10 · commit 64c84e61095e · 2022-05-15T13:01:33.000+05:30
diff --git a/src/libasr/asr_utils.h b/src/libasr/asr_utils.h
@@ -136,6 +136,7 @@ static inline ASR::ttype_t* expr_type(const ASR::expr_t *f)
         case ASR::exprType::ListItem: { return ((ASR::ListItem_t*)f)->m_type; }
         case ASR::exprType::ListSection: { return ((ASR::ListSection_t*)f)->m_type; }
         case ASR::exprType::ListPop: { return ((ASR::ListPop_t*)f)->m_type; }
+        case ASR::exprType::DictPop: { return ((ASR::DictPop_t*)f)->m_type; }
         default : throw LFortranException("Not implemented");
     }
 }
@@ -309,6 +310,7 @@ static inline ASR::expr_t* expr_value(ASR::expr_t *f)
         case ASR::exprType::ListItem: { return ASR::down_cast<ASR::ListItem_t>(f)->m_value; }
         case ASR::exprType::ListSection: { return ASR::down_cast<ASR::ListSection_t>(f)->m_value; }
         case ASR::exprType::ListPop: { return ASR::down_cast<ASR::ListPop_t>(f)->m_value; }
+        case ASR::exprType::DictPop: { return ASR::down_cast<ASR::DictPop_t>(f)->m_value; }
         case ASR::exprType::DictItem: // Drop through
         case ASR::exprType::ArrayConstant: // Drop through
         case ASR::exprType::IntegerConstant: // Drop through
diff --git a/src/lpython/semantics/python_attribute_eval.h b/src/lpython/semantics/python_attribute_eval.h
@@ -24,7 +24,8 @@ struct AttributeHandler {
             {"list@pop", &eval_list_pop},
             {"set@add", &eval_set_add},
             {"set@remove", &eval_set_remove},
-            {"dict@get", &eval_dict_get}
+            {"dict@get", &eval_dict_get},
+            {"dict@pop", &eval_dict_pop}
         };
     }
 
@@ -269,6 +270,29 @@ struct AttributeHandler {
         return make_DictItem_t(al, loc, s, args[0], def, value_type);
     }
 
+    static ASR::asr_t* eval_dict_pop(ASR::symbol_t *s, Allocator &al, const Location &loc,
+            Vec<ASR::expr_t*> &args, diag::Diagnostics &diag) {
+        if (args.size() != 1) {
+            throw SemanticError("'pop' takes only one argument for now", loc);
+        }
+        ASR::Variable_t *v = ASR::down_cast<ASR::Variable_t>(s);
+        ASR::ttype_t *type = v->m_type;
+        ASR::ttype_t *key_type = ASR::down_cast<ASR::Dict_t>(type)->m_key_type;
+        ASR::ttype_t *value_type = ASR::down_cast<ASR::Dict_t>(type)->m_value_type;
+        if (!ASRUtils::check_equal_type(ASRUtils::expr_type(args[0]), key_type)) {
+            std::string ktype = ASRUtils::type_to_str_python(ASRUtils::expr_type(args[0]));
+            std::string totype = ASRUtils::type_to_str_python(key_type);
+            diag.add(diag::Diagnostic(
+                "Type mismatch in pop's key value, the types must be compatible",
+                diag::Level::Error, diag::Stage::Semantic, {
+                    diag::Label("type mismatch (found: '" + ktype + "', expected: '" + totype + "')",
+                            {args[0]->base.loc})
+                })
+            );
+            throw SemanticAbort();
+        }
+        return make_DictPop_t(al, loc, s, args[0], value_type, nullptr);
+    }
 
 }; // AttributeHandler
 
