Add str.capitalize() to ast->asr

AbdelrahmanKhaledd · AbdelrahmanKhaledd · commit 71fb600c20b8 · 2022-08-16T00:48:25.000+02:00
diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp
@@ -4243,6 +4243,23 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                             }
                             ASR::expr_t *se = ASR::down_cast<ASR::expr_t>(
                                             ASR::make_Var_t(al, x.base.base.loc, st));
+                            if (ASR::is_a<ASR::Character_t>(*(ASRUtils::expr_type(se)))) {
+                                if (std::string(at->m_attr) == std::string("capitalize")) {
+                                    if(args.size() != 0) {
+                                        throw SemanticError("str.capitalize() takes no arguments",
+                                            x.base.base.loc);
+                                    }
+                                    ASR::symbol_t *fn_div = resolve_intrinsic_function(x.base.base.loc, "_lpython_str_capitalize");
+                                    Vec<ASR::call_arg_t> args;
+                                    args.reserve(al, 1);
+                                    ASR::call_arg_t arg;
+                                    arg.loc = x.base.base.loc;
+                                    arg.m_value = se;
+                                    args.push_back(al, arg);
+                                    tmp = make_call_helper(al, fn_div, current_scope, args, "_lpython_str_capitalize", x.base.base.loc);
+                                    return;
+                                }
+                            }
                             handle_attribute(se, at->m_attr, x.base.base.loc, eles);
                             return;
                         }
@@ -4276,6 +4293,23 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
                     throw SemanticError("'int' object has no attribute '" + std::string(at->m_attr) + "'",
                         x.base.base.loc);
                 }
+            } else if (AST::is_a<AST::ConstantStr_t>(*at->m_value)) {
+                if (std::string(at->m_attr) == std::string("capitalize")) {
+                    if(args.size() != 0) {
+                        throw SemanticError("str.capitalize() takes no arguments",
+                            x.base.base.loc);
+                    }
+                    AST::ConstantStr_t *n = AST::down_cast<AST::ConstantStr_t>(at->m_value);
+                    std::string res = n->m_value;
+                    res[0] = toupper(res[0]);
+                    ASR::ttype_t *str_type = ASRUtils::TYPE(ASR::make_Character_t(al, x.base.base.loc,
+                                    1, 1, nullptr, nullptr , 0));
+                    tmp = ASR::make_StringConstant_t(al, x.base.base.loc, s2c(al, res), str_type);
+                    return;
+                } else {
+                    throw SemanticError("'str' object has no attribute '" + std::string(at->m_attr) + "'",
+                        x.base.base.loc);
+                }
             } else {
                 throw SemanticError("Only Name type and constant integers supported in Call",
                     x.base.base.loc);
diff --git a/src/lpython/semantics/python_attribute_eval.h b/src/lpython/semantics/python_attribute_eval.h
@@ -6,7 +6,6 @@
 #include <libasr/string_utils.h>
 #include <lpython/utils.h>
 #include <lpython/semantics/semantic_exception.h>
-
 namespace LFortran {
 
 struct AttributeHandler {
@@ -331,6 +330,7 @@ struct AttributeHandler {
 
 }; // AttributeHandler
 
+
 } // namespace LFortran
 
 #endif /* LPYTHON_ATTRIBUTE_EVAL_H */
diff --git a/src/lpython/semantics/python_comptime_eval.h b/src/lpython/semantics/python_comptime_eval.h
@@ -61,7 +61,8 @@ struct PythonIntrinsicProcedures {
             {"_lpython_floordiv", {m_builtin, &eval__lpython_floordiv}},
             {"_mod", {m_builtin, &eval__mod}},
             {"max" , {m_builtin , &eval_max}},
-            {"min" , {m_builtin , &eval_min}}
+            {"min" , {m_builtin , &eval_min}},
+            {"_lpython_str_capitalize", {m_builtin, &eval__lpython_str_capitalize}}
         };
     }
 
@@ -592,8 +593,8 @@ struct PythonIntrinsicProcedures {
         ASR::expr_t *first_element = args[0];
         ASR::ttype_t *first_element_type = ASRUtils::expr_type(first_element);
         semantic_error_flag &= ASRUtils::is_integer(*first_element_type)
-                               || ASRUtils::is_real(*first_element_type)
-                               || ASRUtils::is_character(*first_element_type);
+                                || ASRUtils::is_real(*first_element_type)
+                                || ASRUtils::is_character(*first_element_type);
         int32_t smallest_ind = 0;
         if (semantic_error_flag) {
             if (ASRUtils::is_integer(*first_element_type)) {
@@ -650,6 +651,24 @@ struct PythonIntrinsicProcedures {
 
     }
 
+    static ASR::expr_t *eval__lpython_str_capitalize(Allocator &al, const Location &loc, Vec<ASR::expr_t *> &args) {
+        LFORTRAN_ASSERT(ASRUtils::all_args_evaluated(args));
+        if (args.size() != 1) {
+            throw SemanticError("_lpython_str_capitalize() takes exactly one arguments (" +
+                std::to_string(args.size()) + " given)", loc);
+        }
+        ASR::expr_t *arg1 = args[0];
+        ASR::ttype_t *arg1_type = ASRUtils::expr_type(arg1);
+        if (ASRUtils::is_character(*arg1_type)) {
+            std::string val = ASR::down_cast<ASR::StringConstant_t>(arg1)->m_s;
+            ASR::ttype_t *type = ASRUtils::TYPE(ASR::make_Character_t(al, loc,
+                                        1, 1, nullptr, nullptr, 0));
+            ASR::ttype_t *res_type = ASRUtils::TYPE(ASR::make_StringConstant_t(al, loc, s2c(al, ""), type));
+            return ASR::down_cast<ASR::expr_t>(ASR::make_StringConstant_t(al, loc, s2c(al, val),  res_type));
+        } else {
+            throw SemanticError("Only string from arguments.", loc);
+        }
+    }
 
 
 }; // ComptimeEval
