Repository: pyccel/pyccel (python)

Issue:
Add support for strings as arguments
### Describe the feature
It should be possible to pass strings as arguments to functions

### Test Code
Provide code which does not currently work but which should do when this issue is fixed:
```python
def f(option : str):
    if option == 'do this':
        return 1.0
    else:
        return 2.0
```

### Proposed Solution
Most of the tools necessary for this should already be implemented. The main missing piece is the wrapping

### Additional context
This feature was requested by @spossann for Struphy

Pull request:
Allow passing strings as arguments
Add the necessary wrappers to add support for strings as arguments. Fixes #2194 .
Add support for the built-in function `str` so it can be used as a type annotation.
Correct the type of `s[0]` when `s` is a string (should be string, previously char).
Add basic tests.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d349ec33..4e33e284 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -73,6 +73,7 @@ All notable changes to this project will be documented in this file.
 -   #2183 : Add compile time analysis of if block conditions.
 -   #2139 : Add support for `__getitem__`
 -   #337 : Add support for returning tuples from functions.
+-   #2194 : Add support for strings as arguments.
 -   \[INTERNALS\] Add abstract class `SetMethod` to handle calls to various set methods.
 -   \[INTERNALS\] Added `container_rank` property to `ast.datatypes.PyccelType` objects.
 -   \[INTERNALS\] Add a `__call__` method to `FunctionDef` to create `FunctionCall` instances.
diff --git a/docs/builtin-functions.md b/docs/builtin-functions.md
index 7ddd157c..6c8b6823 100644
--- a/docs/builtin-functions.md
+++ b/docs/builtin-functions.md
@@ -65,7 +65,7 @@ Python contains a limited number of builtin functions defined [here](https://doc
 | `slice` | No |
 | `sorted` | No |
 | `staticmethod` | No |
-| `str` | No |
+| `str` | As a type annotation or with string arguments |
 | **`sum`** | **Yes** |
 | `super` | No |
 | **`tuple`** | **Yes** |
diff --git a/docs/type_annotations.md b/docs/type_annotations.md
index b23840ff..f902c0d5 100644
--- a/docs/type_annotations.md
+++ b/docs/type_annotations.md
@@ -105,9 +105,12 @@ Strings are not yet supported as keys in Fortran.
 
 ## Strings
 
-Pyccel contains very minimal support for strings. This is mostly provided to allow the use of strings as keys of dictionaries.
+Pyccel contains very minimal support for strings. For example strings can be used as keys of dictionaries (only in C currently) or to represent flags in the user code with if statements checking their value. More complex string handling is not currently supported. See the documentation on [builtin functions](./builtin-functions.md) for an overview of the available string methods.
 
-Currently strings can be used as local variables and can be returned from functions.
+Strings can be declared in Pyccel using the following syntax:
+```python
+a : str = 'hello'
+```
 
 ## Handling multiple types
 
diff --git a/pyccel/ast/builtins.py b/pyccel/ast/builtins.py
index f57d27ba..f6e5bd0c 100644
--- a/pyccel/ast/builtins.py
+++ b/pyccel/ast/builtins.py
@@ -16,7 +16,7 @@ from pyccel.utilities.stage import PyccelStage
 
 from .basic     import PyccelAstNode, TypedAstNode
 from .datatypes import PythonNativeInt, PythonNativeBool, PythonNativeFloat
-from .datatypes import GenericType, PythonNativeComplex
+from .datatypes import GenericType, PythonNativeComplex, CharType
 from .datatypes import PrimitiveBooleanType, PrimitiveComplexType
 from .datatypes import HomogeneousTupleType, InhomogeneousTupleType, TupleType
 from .datatypes import HomogeneousListType, HomogeneousContainerType
@@ -58,6 +58,7 @@ __all__ = (
     'PythonRound',
     'PythonSet',
     'PythonSetFunction',
+    'PythonStr',
     'PythonSum',
     'PythonTuple',
     'PythonTupleFunction',
@@ -1724,6 +1725,36 @@ class PythonIsInstance(PyccelFunction):
     def __init__(self, obj, class_or_tuple):
         super().__init__(obj, class_or_tuple)
 
+#==============================================================================
+class PythonStr(PyccelFunction):
+    """
+    Represents a call to Python's `str` function.
+
+    Represents a call to Python's `str` function which describes a string
+    cast.
+
+    Parameters
+    ----------
+    arg : TypedAstNode
+        The argument that is cast to a string.
+    """
+    __slots__ = ('_shape',)
+    _static_type = StringType()
+    _class_type = StringType()
+    name = 'str'
+
+    def __new__(cls, arg):
+        if isinstance(arg, LiteralString):
+            return arg
+        else:
+            return super().__new__(cls)
+
+    def __init__(self, arg):
+        if not isinstance(arg.class_type, (StringType, CharType)):
+            raise NotImplementedError("Support for casting non-character types to strings is not yet available")
+        self._shape = (None,)
+        super().__init__(arg)
+
 #==============================================================================
 
 DtypePrecisionToCastFunction = {
@@ -1753,7 +1784,7 @@ builtin_functions_dict = {
     'range'      : PythonRange,
     'round'      : PythonRound,
     'set'        : PythonSetFunction,
-    'str'        : LiteralString,
+    'str'        : PythonStr,
     'sum'        : PythonSum,
     'tuple'      : PythonTupleFunction,
     'type'       : PythonType,
diff --git a/pyccel/ast/cwrapper.py b/pyccel/ast/cwrapper.py
index 6c7c193e..c672ebc7 100644
--- a/pyccel/ast/cwrapper.py
+++ b/pyccel/ast/cwrapper.py
@@ -1281,6 +1281,27 @@ PyDict_SetItem = FunctionDef(name = 'PyDict_SetItem',
                     body = [])
 
 
+#-------------------------------------------------------------------
+#                         String functions
+#-------------------------------------------------------------------
+# https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_AsUTF8
+PyUnicode_AsUTF8 = FunctionDef(name = 'PyUnicode_AsUTF8',
+                    arguments = [FunctionDefArgument(Variable(PyccelPyObject(), 'unicode', memory_handling='alias'))],
+                    results = FunctionDefResult(Variable(CharType(), 'str', memory_handling='alias')),
+                    body = [])
+
+# https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_Check
+PyUnicode_Check = FunctionDef(name = 'PyUnicode_Check',
+                    arguments = [FunctionDefArgument(Variable(PyccelPyObject(), 'str', memory_handling='alias'))],
+                    results = FunctionDefResult(Variable(CharType(), 'out', memory_handling='alias')),
+                    body = [])
+
+# https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_GetLength
+PyUnicode_GetLength = FunctionDef(name = 'PyUnicode_GetLength',
+                    arguments = [FunctionDefArgument(Variable(PyccelPyObject(), 'str', memory_handling='alias'))],
+                    results = FunctionDefResult(Variable(PythonNativeInt(), 'len')),
+                    body = [])
+
 # Functions definitions are defined in pyccel/stdlib/cwrapper/cwrapper.c
 check_type_registry = {
     PythonNativeBool()    : 'PyIs_Bool',
diff --git a/pyccel/ast/literals.py b/pyccel/ast/literals.py
index fcf292eb..b95967c1 100644
--- a/pyccel/ast/literals.py
+++ b/pyccel/ast/literals.py
@@ -331,7 +331,6 @@ class LiteralString(Literal):
     """
     __slots__ = ('_string',)
     _class_type = StringType()
-    _static_type = StringType()
     _shape = (None,)
 
     def __init__(self, arg):
diff --git a/pyccel/ast/variable.py b/pyccel/ast/variable.py
index 14f3f728..768e9f75 100644
--- a/pyccel/ast/variable.py
+++ b/pyccel/ast/variable.py
@@ -14,7 +14,7 @@ from pyccel.utilities.stage import PyccelStage
 
 from .basic     import PyccelAstNode, TypedAstNode
 from .datatypes import PyccelType, InhomogeneousTupleType, HomogeneousListType, HomogeneousSetType, DictType
-from .datatypes import ContainerType, HomogeneousTupleType
+from .datatypes import ContainerType, HomogeneousTupleType, CharType, StringType
 from .internals import PyccelArrayShapeElement, Slice, PyccelSymbol
 from .internals import apply_pickle
 from .literals  import LiteralInteger, Nil, LiteralEllipsis
@@ -731,6 +731,10 @@ class IndexedElement(TypedAstNode):
                 self._is_slice = True
                 self._shape = tuple(new_shape)
 
+            if isinstance(self._class_type, CharType):
+                self._class_type = StringType()
+                self._shape = (1,)
+
         if isinstance(self._class_type, InhomogeneousTupleType) and self._shape is None:
             self._shape = (LiteralInteger(len(self._class_type)),)
 
diff --git a/pyccel/codegen/printing/ccode.py b/pyccel/codegen/printing/ccode.py
index 456603eb..b0ac86c3 100644
--- a/pyccel/codegen/printing/ccode.py
+++ b/pyccel/codegen/printing/ccode.py
@@ -3032,6 +3032,8 @@ class CCodePrinter(CodePrinter):
 
         return f"(*{container_type}_at({dict_obj_code}, {key}))"
 
+    #================== String methods ==================
+
     def _print_CStrStr(self, expr):
         arg = expr.args[0]
         code = self._print(ObjectAddress(arg))
@@ -3040,6 +3042,15 @@ class CCodePrinter(CodePrinter):
         else:
             return f'cstr_str({code})'
 
+    def _print_PythonStr(self, expr):
+        arg = expr.args[0]
+        arg_code = self._print(arg)
+        if isinstance(arg.class_type, StringType):
+            return f'cstr_clone({arg_code})'
+        else:
+            assert isinstance(arg.class_type, CharType) and getattr(arg, 'is_alias', True)
+            return f'cstr_from({arg_code})'
+
     #=================== MACROS ==================
 
     def _print_MacroShape(self, expr):
diff --git a/pyccel/codegen/printing/fcode.py b/pyccel/codegen/printing/fcode.py
index ae613d48..5e1c43d9 100644
--- a/pyccel/codegen/printing/fcode.py
+++ b/pyccel/codegen/printing/fcode.py
@@ -1495,7 +1495,7 @@ class FCodePrinter(CodePrinter):
         # See pyccel/stdlib/gFTL_functions/Set_extensions.inc for the definition
         return f'{type_name}_is_disjoint({var_code}, {arg_code})'
 
-   #========================== Dict Methods ================================#
+    #========================== Dict Methods ================================#
 
     def _print_DictClear(self, expr):
         var = self._print(expr.dict_obj)
@@ -1506,6 +1506,11 @@ class FCodePrinter(CodePrinter):
         key = self._print(expr.key)
         return f'{dict_obj} % of( {key} )'
 
+    #========================== String Methods ===============================#
+
+    def _print_PythonStr(self, expr):
+        return self._print(expr.args[0])
+
     #========================== Numpy Elements ===============================#
 
     def _print_NumpySum(self, expr):
@@ -2592,7 +2597,7 @@ class FCodePrinter(CodePrinter):
         body_code = self._print(expr.body)
         docstring = self._print(expr.docstring) if expr.docstring else ''
 
-        decs = [Declare(v) for v in expr.local_vars]
+        decs = [Declare(v) for v in expr.local_vars if not v.is_argument]
         self._get_external_declarations(decs)
 
         prelude += ''.join(self._print(i) for i in decs)
@@ -3128,7 +3133,7 @@ class FCodePrinter(CodePrinter):
 
     def _print_PyccelAdd(self, expr):
         if isinstance(expr.dtype, StringType):
-            return '//'.join('trim('+self._print(a)+')' for a in expr.args)
+            return ' // '.join(self._print(a) for a in expr.args)
         else:
             args = [PythonInt(a) if isinstance(a.dtype.primitive_type, PrimitiveBooleanType) else a for a in expr.args]
             return ' + '.join(self._print(a) for a in args)
diff --git a/pyccel/codegen/wrapper/c_to_python_wrapper.py b/pyccel/codegen/wrapper/c_to_python_wrapper.py
index 23e49566..c20527ea 100644
--- a/pyccel/codegen/wrapper/c_to_python_wrapper.py
+++ b/pyccel/codegen/wrapper/c_to_python_wrapper.py
@@ -12,7 +12,7 @@ from pyccel.ast.bind_c        import BindCFunctionDef, BindCPointer
 from pyccel.ast.bind_c        import BindCModule, BindCModuleVariable, BindCVariable
 from pyccel.ast.bind_c        import BindCClassDef, BindCClassProperty, BindCArrayType
 from pyccel.ast.builtins      import PythonTuple, PythonRange, PythonLen, PythonSet
-from pyccel.ast.builtins      import VariableIterator
+from pyccel.ast.builtins      import VariableIterator, PythonStr
 from pyccel.ast.builtin_methods.set_methods import SetAdd, SetPop
 from pyccel.ast.builtin_methods.dict_methods import DictItems
 from pyccel.ast.class_defs    import StackArrayClass
@@ -38,6 +38,7 @@ from pyccel.ast.cwrapper      import PySet_New, PySet_Add
 from pyccel.ast.cwrapper      import PySet_Size, PySet_Check, PySet_GetIter, PySet_Clear
 from pyccel.ast.cwrapper      import PyIter_Next
 from pyccel.ast.cwrapper      import PyDict_New, PyDict_SetItem
+from pyccel.ast.cwrapper      import PyUnicode_AsUTF8, PyUnicode_Check, PyUnicode_GetLength
 from pyccel.ast.c_concepts    import ObjectAddress, PointerCast, CStackArray, CNativeInt
 from pyccel.ast.c_concepts    import CStrStr
 from pyccel.ast.datatypes     import VoidType, PythonNativeInt, CustomDataType, DataTypeFactory
@@ -310,6 +311,8 @@ class CToPythonWrapper(Wrapper):
         if isinstance(dtype, CustomDataType):
             python_cls_base = self.scope.find(dtype.name, 'classes', raise_if_missing = True)
             type_check_condition = PyObject_TypeCheck(py_obj, python_cls_base.type_object)
+        elif isinstance(dtype, StringType):
+            type_check_condition = PyUnicode_Check(py_obj)
         elif rank == 0:
             try :
                 cast_function = check_type_registry[dtype]
@@ -2580,6 +2583,87 @@ class CToPythonWrapper(Wrapper):
 
         return {'body': body, 'args': arg_vars, 'clean_up': clean_up}
 
+    def _extract_StringType_FunctionDefArgument(self, orig_var, collect_arg, bound_argument,
+            is_bind_c_argument, *, arg_var = None):
+        """
+        Extract the C-compatible string FunctionDefArgument from the PythonObject.
+
+        Extract the C-compatible string FunctionDefArgument from the PythonObject.
+        The C-compatible argument is extracted from collect_arg which holds a Python
+        object into arg_var.
+
+        The extraction is done by allocating an array and filling the elements with values
+        extracted from the indexed Python tuple in collect_arg.
+
+        Parameters
+        ----------
+        orig_var : Variable | IndexedElement
+            An object representing the variable or an element of the variable from the
+            FunctionDefArgument being wrapped.
+
+        collect_arg : Variable
+            A variable with type PythonObject* holding the Python argument from which the
+            C-compatible argument should be collected.
+
+        bound_argument : bool
+            True if the argument is the self argument of a class method. False otherwise.
+            This should always be False for this function.
+
+        is_bind_c_argument : bool
+            True if the argument was saved in a BindCFunctionDefArgument. False otherwise.
+
+        arg_var : Variable | IndexedElement, optional
+            A variable or an element of the variable representing the argument that
+            will be passed to the low-level function call.
+
+        Returns
+        -------
+        list[PyccelAstNode]
+            A list of expressions which extract the argument from collect_arg into arg_var.
+        """
+        assert bound_argument is False
+
+        if is_bind_c_argument:
+            if arg_var is None:
+                data_var = Variable(CStackArray(CharType()), self.scope.get_expected_name(orig_var.name),
+                                    shape = (None,), memory_handling='alias', is_const = True)
+                size_var = Variable(PythonNativeInt(), self.scope.get_new_name(f'{data_var.name}_size'))
+                arg_var = Variable(BindCArrayType(1, False), self.scope.get_new_name(orig_var.name),
+                                    shape = (LiteralInteger(2),))
+                self.scope.insert_variable(data_var, orig_var.name)
+                self.scope.insert_variable(size_var)
+                self.scope.insert_variable(arg_var, tuple_recursive = False)
+                self.scope.insert_symbolic_alias(arg_var[0], ObjectAddress(data_var))
+                self.scope.insert_symbolic_alias(arg_var[1], size_var)
+
+            if getattr(orig_var, 'is_optional', False):
+                body = [AliasAssign(orig_var, PyUnicode_AsUTF8(collect_arg)),
+                        Assign(self.scope.collect_tuple_element(arg_var[1]), PyUnicode_GetLength(collect_arg))]
+            else:
+                body = [Assign(orig_var, PyUnicode_AsUTF8(collect_arg)),
+                        Assign(self.scope.collect_tuple_element(arg_var[1]), PyUnicode_GetLength(collect_arg))]
+
+            default_init = [AliasAssign(data_var, Nil()),
+                            Assign(size_var, 0)]
+        else:
+
+            if arg_var is None:
+                arg_var = orig_var.clone(self.scope.get_expected_name(orig_var.name), new_class = Variable,
+                                        is_argument = False)
+                self.scope.insert_variable(arg_var, orig_var.name)
+
+            body = [Assign(orig_var, PythonStr(PyUnicode_AsUTF8(collect_arg)))]
+
+            default_init = [AliasAssign(arg_var, Nil())]
+            if getattr(orig_var, 'is_optional', False):
+                memory_var = self.scope.get_temporary_variable(arg_var, name = arg_var.name + '_memory', is_optional = False,
+                                                clone_scope = self.scope)
+                body.insert(0, AliasAssign(arg_var, memory_var))
+
+        return {'body': body,
+                'args': [arg_var],
+                'default_init': default_init}
+
     def _extract_FunctionDefResult(self, orig_var, is_bind_c, funcdef = None):
         """
         Get the code which translates a C-compatible `Variable` to a Python `FunctionDefResult`.
@@ -2978,13 +3062,15 @@ class CToPythonWrapper(Wrapper):
             c_res = Variable(CharType(), self.scope.get_new_name(name+'_data'), memory_handling='alias')
             self.scope.insert_variable(c_res)
             char_data = ObjectAddress(c_res)
+            result = [char_data]
         else:
             c_res = Variable(StringType(), self.scope.get_new_name(name), memory_handling='heap')
             self.scope.insert_variable(c_res)
             char_data = CStrStr(c_res)
+            result = [c_res]
 
         body = [AliasAssign(py_res, PyBuildValueNode([char_data]))]
         if is_bind_c:
             body.append(Deallocate(c_res))
-        return {'c_results': [c_res], 'py_result': py_res, 'body': body}
+        return {'c_results': result, 'py_result': py_res, 'body': body}
 
diff --git a/pyccel/codegen/wrapper/fortran_to_c_wrapper.py b/pyccel/codegen/wrapper/fortran_to_c_wrapper.py
index f1dc9e52..bd9423e8 100644
--- a/pyccel/codegen/wrapper/fortran_to_c_wrapper.py
+++ b/pyccel/codegen/wrapper/fortran_to_c_wrapper.py
@@ -26,7 +26,7 @@ from pyccel.ast.datatypes import HomogeneousTupleType, TupleType
 from pyccel.ast.datatypes import PythonNativeInt, CharType
 from pyccel.ast.datatypes import InhomogeneousTupleType
 from pyccel.ast.internals import Slice
-from pyccel.ast.literals import LiteralInteger, Nil, LiteralTrue
+from pyccel.ast.literals import LiteralInteger, Nil, LiteralTrue, LiteralString
 from pyccel.ast.numpytypes import NumpyNDArrayType, NumpyInt32Type
 from pyccel.ast.operators import PyccelIsNot, PyccelMul, PyccelAdd
 from pyccel.ast.variable import Variable, IndexedElement, DottedVariable
@@ -407,6 +407,49 @@ class FortranToCWrapper(Wrapper):
 
         return {'c_arg': BindCVariable(c_arg_var, var), 'f_arg': arg_var, 'body': body}
 
+    def _extract_StringType_FunctionDefArgument(self, var, func):
+        name = var.name
+        scope = self.scope
+        scope.insert_symbol(name)
+        collisionless_name = scope.get_expected_name(name)
+        rank = var.rank
+        bind_var = Variable(BindCPointer(), scope.get_new_name(f'bound_{name}'),
+                            is_argument = True, is_optional = False, memory_handling='alias',
+                            is_const = True)
+        arg_var = var.clone(collisionless_name, is_argument = False, is_optional = False,
+                            allows_negative_indexes=False, new_class = Variable)
+        array_var = Variable(NumpyNDArrayType(CharType(), 1, None), scope.get_new_name(name),
+                            memory_handling='alias')
+        scope.insert_variable(arg_var)
+        scope.insert_variable(bind_var)
+        scope.insert_variable(array_var)
+
+        shape_var = scope.get_temporary_variable(PythonNativeInt(), name=f'{name}_size', is_argument = True)
+
+        for_scope = scope.create_new_loop_scope()
+        iterator = PythonRange(LiteralInteger(1), PyccelAdd(shape_var, LiteralInteger(1)))
+        idx = Variable(PythonNativeInt(), self.scope.get_new_name())
+        iterator.set_loop_counter(idx)
+        self.scope.insert_variable(idx)
+
+        # Default Fortran arrays retrieved from C_F_Pointer are 1-indexed
+        # Lists are 1-indexed but Pyccel adds the shift during printing so they are
+        # treated as 0-indexed here
+        for_body = [Assign(arg_var, PyccelAdd(arg_var, IndexedElement(array_var, idx)))]
+
+        body = [C_F_Pointer(bind_var, array_var, (shape_var,)),
+                Assign(arg_var, LiteralString('')),
+                For((idx,), iterator, for_body, scope = for_scope)]
+
+        c_arg_var = Variable(BindCArrayType(rank, has_strides = False),
+                        scope.get_new_name(), is_argument = True,
+                        shape = (LiteralInteger(2),))
+
+        scope.insert_symbolic_alias(IndexedElement(c_arg_var, LiteralInteger(0)), bind_var)
+        scope.insert_symbolic_alias(IndexedElement(c_arg_var, LiteralInteger(1)), shape_var)
+
+        return {'c_arg': BindCVariable(c_arg_var, var), 'f_arg': arg_var, 'body': body}
+
     def _extract_HomogeneousSetType_FunctionDefArgument(self, var, func):
         name = var.name
         scope = self.scope


Test patch:
diff --git a/tests/epyccel/modules/strings_module.py b/tests/epyccel/modules/strings_module.py
index 06ddcf41..9659344a 100644
--- a/tests/epyccel/modules/strings_module.py
+++ b/tests/epyccel/modules/strings_module.py
@@ -8,6 +8,9 @@ __all__ = [
         'concatenate',
         'concatenate_multiple',
         'concatenate_expr',
+        'string_function_call',
+        'string_function_call_on_literal',
+        'string_function_return',
         ]
 
 def one_quote():
@@ -47,3 +50,15 @@ def concatenate_expr():
     t = "world"
     v = s + "_" + t
     return v
+
+def string_function_call():
+    s = 'hello'
+    t = str(s)
+    return t
+
+def string_function_call_on_literal():
+    t = str('hello')
+    return t
+
+def string_function_return():
+    return str('hello')
diff --git a/tests/epyccel/test_strings.py b/tests/epyccel/test_strings.py
index 676c3da2..955024eb 100644
--- a/tests/epyccel/test_strings.py
+++ b/tests/epyccel/test_strings.py
@@ -42,3 +42,29 @@ def test_string_compare(language):
     f = epyccel( str_comp, language=language )
 
     assert str_comp() == f()
+
+def test_string_argument(language):
+    def str_option_test(option : str):
+        if option == 'do this':
+            return 1.0
+        else:
+            return 2.0
+
+    f = epyccel( str_option_test, language=language )
+
+    assert str_option_test('do this') == f('do this')
+    assert str_option_test('do that') == f('do that')
+
+def test_string_argument_optional(language):
+    def str_option_test(option : str = None):
+        if option is not None and option == 'do this':
+            return 1.0
+        else:
+            return 2.0
+
+    f = epyccel( str_option_test, language=language )
+
+    assert str_option_test('do this') == f('do this')
+    assert str_option_test('do that') == f('do that')
+    assert str_option_test() == f()
+