diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt index 4768d42514..f01a218a36 100644 --- a/integration_tests/CMakeLists.txt +++ b/integration_tests/CMakeLists.txt @@ -619,5 +619,6 @@ RUN(NAME intrinsics_01 LABELS cpython llvm NOFAST) # any # lpython decorator RUN(NAME lpython_decorator_01 LABELS cpython) +RUN(NAME lpython_decorator_02 LABELS cpython) COMPILE(NAME import_order_01 LABELS cpython llvm c) # any diff --git a/integration_tests/lpython_decorator_02.py b/integration_tests/lpython_decorator_02.py new file mode 100644 index 0000000000..3c64d7a640 --- /dev/null +++ b/integration_tests/lpython_decorator_02.py @@ -0,0 +1,19 @@ +from numpy import array +from lpython import i32, f64, lpython, TypeVar + +n = TypeVar("n") + +@lpython +def multiply(n: i32, x: f64[:]) -> f64[n]: + i: i32 + for i in range(n): + x[i] *= 5.0 + return x + +def test(): + size: i32 = 5 + x: f64[5] = array([11.0, 12.0, 13.0, 14.0, 15.0]) + y: f64[5] = (multiply(size, x)) + assert y[2] == 65. + +test() diff --git a/src/runtime/lpython/lpython.py b/src/runtime/lpython/lpython.py index ce769d1165..b38d214142 100644 --- a/src/runtime/lpython/lpython.py +++ b/src/runtime/lpython/lpython.py @@ -637,7 +637,8 @@ def get_rtlib_dir(): def get_type_info(arg): # return_type -> (`type_format`, `variable type`, `array struct name`) - # See: https://docs.python.org/3/c-api/arg.html for more info on type_format + # See: https://docs.python.org/3/c-api/arg.html for more info on `type_format` + # `array struct name`: used by the C backend if arg == f64: return ('d', "double", 'r64') elif arg == f32: @@ -652,7 +653,10 @@ def get_type_info(arg): t = get_type_info(arg._type) if t[2] == '': raise NotImplementedError("Type %r not implemented" % arg) - return ('O', ["PyArrayObject *", "struct "+t[2]+" *", t[1]+" *"], '') + n = '' + if not isinstance(arg._dims, slice): + n = arg._dims._name + return ('O', ["PyArrayObject *", "struct "+t[2]+" *", t[1]+" *", n], '') else: raise NotImplementedError("Type %r not implemented" % arg) @@ -662,6 +666,18 @@ def get_data_type(t): else: return t + " " + def get_typenum(t): + if t == "int": + return "NPY_INT" + elif t == "long int": + return "NPY_LONG" + elif t == "float": + return "NPY_FLOAT" + elif t == "double": + return "NPY_DOUBLE" + else: + raise NotImplementedError("Type %s not implemented" % t) + self.fn_name = function.__name__ # Get the source code of the function source_code = getsource(function) @@ -682,29 +698,36 @@ def get_data_type(t): self.arg_type_formats = "" self.return_type = "" self.return_type_format = "" + self.array_as_return_type = () self.arg_types = {} - counter = 1 for t in types.keys(): if t == "return": type = get_type_info(types[t]) - self.return_type_format = type[0] - self.return_type = type[1] + if type[0] == 'O': + self.array_as_return_type = type + continue + else: + self.return_type_format = type[0] + self.return_type = type[1] else: type = get_type_info(types[t]) self.arg_type_formats += type[0] - self.arg_types[counter] = type[1] - counter += 1 + self.arg_types[t] = type[1] # ---------------------------------------------------------------------- - # `arg_0`: used as the return variables - # arguments are declared as `arg_1`, `arg_2`, ... - variables_decl = "" + # `__return_value`: used as the return variables + variables_decl = "// Declare return variables and arguments\n" if self.return_type != "": - variables_decl = "// Declare return variables and arguments\n" - variables_decl += " " + get_data_type(self.return_type) + "arg_" \ - + str(0) + ";\n" + variables_decl += " " + get_data_type(self.return_type) \ + + "_" + self.fn_name + "_return_value;\n" + elif self.array_as_return_type: + variables_decl += " " + self.array_as_return_type[1][1] + "_" \ + + self.fn_name + "_return_value = malloc(sizeof(" \ + + self.array_as_return_type[1][1][:-2] + "));\n" + else: + variables_decl = "" # ---------------------------------------------------------------------- # `PyArray_AsCArray` is used to convert NumPy Arrays to C Arrays - # `fill_array_details` contains arrays operations to be + # `fill_array_details` contains array operations to be # performed on the arguments # `parse_args` are used to capture the args from CPython # `pass_args` are the args that are passed to the shared library function @@ -712,16 +735,18 @@ def get_data_type(t): parse_args = "" pass_args = "" numpy_init = "" + prefix_comma = False for i, t in self.arg_types.items(): - if i > 1: + if prefix_comma: parse_args += ", " pass_args += ", " + prefix_comma = True if isinstance(t, list): if numpy_init == "": numpy_init = "// Initialize NumPy\n import_array();\n\n " fill_array_details += f"""\n - // fill array details for args[{i-1}] - if (PyArray_NDIM(arg_{i}) != 1) {{ + // fill array details for {i} + if (PyArray_NDIM({i}) != 1) {{ PyErr_SetString(PyExc_TypeError, "Only 1 dimension is implemented for now."); return NULL; @@ -731,9 +756,9 @@ def get_data_type(t): {{ {t[2]}array; // Create C arrays from numpy objects: - PyArray_Descr *descr = PyArray_DescrFromType(PyArray_TYPE(arg_{i})); + PyArray_Descr *descr = PyArray_DescrFromType(PyArray_TYPE({i})); npy_intp dims[1]; - if (PyArray_AsCArray((PyObject **)&arg_{i}, (void *)&array, dims, 1, descr) < 0) {{ + if (PyArray_AsCArray((PyObject **)&{i}, (void *)&array, dims, 1, descr) < 0) {{ PyErr_SetString(PyExc_TypeError, "error converting to c array"); return NULL; }} @@ -744,11 +769,11 @@ def get_data_type(t): s_array_{i}->dims[0].length = dims[0]; s_array_{i}->is_allocated = false; }}""" - pass_args += "s_array_" + str(i) + pass_args += "s_array_" + i else: - pass_args += "arg_" + str(i) - variables_decl += " " + get_data_type(t) + "arg_" + str(i) + ";\n" - parse_args += "&arg_" + str(i) + pass_args += i + variables_decl += " " + get_data_type(t) + i + ";\n" + parse_args += "&" + i if parse_args != "": parse_args = f"""\n // Parse the arguments from Python @@ -761,12 +786,38 @@ def get_data_type(t): fill_return_details = "" if self.return_type != "": fill_return_details = f"""\n\n // Call the C function - arg_0 = {self.fn_name}({pass_args}); + _{self.fn_name}_return_value = {self.fn_name}({pass_args}); // Build and return the result as a Python object - return Py_BuildValue("{self.return_type_format}", arg_0);""" + return Py_BuildValue("{self.return_type_format}", _{self.fn_name}_return_value);""" else: - fill_return_details = f"""{self.fn_name}({pass_args}); + if self.array_as_return_type: + fill_return_details = f"""\n + _{self.fn_name}_return_value->data = malloc({self.array_as_return_type[1][3] + } * sizeof({self.array_as_return_type[1][2][:-2]})); + _{self.fn_name}_return_value->n_dims = 1; + _{self.fn_name}_return_value->dims[0].lower_bound = 0; + _{self.fn_name}_return_value->dims[0].length = { + self.array_as_return_type[1][3]}; + _{self.fn_name}_return_value->is_allocated = false; + + // Call the C function + {self.fn_name}({pass_args}, &_{self.fn_name}_return_value[0]); + + // Build and return the result as a Python object + {{ + npy_intp dims[] = {{{self.array_as_return_type[1][3]}}}; + PyObject* numpy_array = PyArray_SimpleNewFromData(1, dims, { + get_typenum(self.array_as_return_type[1][2][:-2])}, + _{self.fn_name}_return_value->data); + if (numpy_array == NULL) {{ + PyErr_SetString(PyExc_TypeError, "error creating an array"); + return NULL; + }} + return numpy_array; + }}""" + else: + fill_return_details = f"""{self.fn_name}({pass_args}); Py_RETURN_NONE;""" # ----------------------------------------------------------------------