Repository: pyccel/pyccel
Pull request #2164 (2025-01-17 13:09:01)

Add support for `dict.__getitem__`
Add support for `dict.__getitem__`. This is important for easy use of dictionaries and will provide some of the work necessary for #2139 .

pyccel-bot[bot]: Hello again! Thank you for this new pull request :star_struck:.

Please begin by requesting your checklist using the command `/bot checklist`

EmilyBourne: /bot checklist
pyccel-bot[bot]: Unfortunately your PR is not passing the tests so it is not quite ready for review yet. Let me know when it is fixed with `/bot mark as ready`.

EmilyBourne: /bot run pylint
EmilyBourne: /bot run docs
pyccel-bot[bot]: Hey @pyccel/pyccel-dev ! @EmilyBourne has just created this great new pull request! Check it out and let me know what you think!

EmilyBourne: /bot run pr_tests
pyccel-bot[bot]: Hey @yguclu, @EmilyBourne, this PR is looking pretty good. @EmilyBourne and @jalalium think it is ready to merge. Could you add your expertise to confirm that this follows all the coding conventions and fits in Pyccel's future plans? Thanks :smile:


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 16f4b90f..2339ad1f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,8 +4,6 @@ All notable changes to this project will be documented in this file.
 ## \[UNRELEASED\]
 
 ### Added
--   #1881 : Add Python support for dict method `copy()`.
--   #1880 : Add support for dict method `clear()`.
 -   #1720 : Add support for `Ellipsis` as the only index for an array.
 -   #1787 : Ensure STC is installed with Pyccel.
 -   #1656 : Ensure gFTL is installed with Pyccel.
@@ -16,7 +14,7 @@ All notable changes to this project will be documented in this file.
 -   #1893 : Add Python support for set initialisation with `set()`.
 -   #1895 : Add Python support for dict initialisation with `{}`.
 -   #1895 : Add Python support for dict initialisation with `dict()`.
--   #1886 : Add Python and C support for dict method `pop()`.
+-   #1881 : Add Python support for dict method `copy()`.
 -   #1887 : Add Python support for dict method `popitem()`.
 -   #1888 : Add Python support for dict method `setdefault()`.
 -   #1885 : Add Python and C support for dict method `get()`.
@@ -49,7 +47,10 @@ All notable changes to this project will be documented in this file.
 -   #1744 : Add support for set method `intersection()`.
 -   #1745 : Add support for set method `intersection_update()`.
 -   #2059 : Add C support for returning dictionaries from functions.
+-   #2164 : Add support for dict indexing.
 -   #1884 : Add support for dict method `items()`.
+-   #1880 : Add support for dict method `clear()`.
+-   #1886 : Add Python and C support for dict method `pop()`.
 -   #1936 : Add missing C output for inline decorator example in documentation
 -   #1937 : Optimise `pyccel.ast.basic.PyccelAstNode.substitute` method.
 -   #1544 : Add support for `typing.TypeAlias`.
diff --git a/pyccel/ast/builtin_methods/dict_methods.py b/pyccel/ast/builtin_methods/dict_methods.py
index dd8a9607..b2ac5d8f 100644
--- a/pyccel/ast/builtin_methods/dict_methods.py
+++ b/pyccel/ast/builtin_methods/dict_methods.py
@@ -11,13 +11,15 @@ This module contains objects which describe these methods within Pyccel's AST.
 """
 
 from pyccel.ast.datatypes import InhomogeneousTupleType, VoidType, SymbolicType
-from pyccel.ast.internals import PyccelFunction, Iterable
+from pyccel.ast.internals import PyccelFunction, Iterable, PyccelArrayShapeElement
+from pyccel.ast.literals  import LiteralInteger
 from pyccel.ast.variable  import IndexedElement
 
 
 __all__ = ('DictClear',
            'DictCopy',
            'DictGet',
+           'DictGetItem',
            'DictItems',
            'DictMethod',
            'DictPop',
@@ -150,13 +152,13 @@ class DictGet(DictMethod):
         The value that should be returned if the key is not present in the
         dictionary.
     """
-    __slots__ = ('_class_type',)
-    _shape = None
+    __slots__ = ('_class_type', '_shape')
     name = 'get'
 
     def __init__(self, dict_obj, k, d = None):
         dict_type = dict_obj.class_type
         self._class_type = dict_type.value_type
+        self._shape = None
         if k.class_type != dict_type.key_type:
             raise TypeError(f"Key passed to get method has type {k.class_type}. Expected {dict_type.key_type}")
         if d and d.class_type != dict_type.value_type:
@@ -164,6 +166,10 @@ class DictGet(DictMethod):
 
         super().__init__(dict_obj, k, d)
 
+        if self._class_type.rank:
+            self._shape = [PyccelArrayShapeElement(self,LiteralInteger(i)) \
+                    for i in range(self._class_type.rank)]
+
     @property
     def key(self):
         """
@@ -213,7 +219,7 @@ class DictSetDefault(DictMethod):
         dict_type = dict_obj.class_type
         self._class_type = dict_type.value_type
 
-        self._shape = (None) * self._class_type.rank if self._class_type.rank else None
+        self._shape = (None,) * self._class_type.rank if self._class_type.rank else None
 
         if k.class_type != dict_type.key_type:
             raise TypeError(f"Key passed to setdefault method has type {k.class_type}. Expected {dict_type.key_type}")
@@ -273,13 +279,13 @@ class DictCopy(DictMethod):
     dict_obj : TypedAstNode
         The object from which the method is called.
     """
-    __slots__ = ('_class_type',)
-    _shape = None
+    __slots__ = ('_class_type', '_shape')
     name = 'copy'
 
     def __init__(self, dict_obj):
         dict_type = dict_obj.class_type
         self._class_type = dict_type
+        self._shape = dict_obj.shape
         super().__init__(dict_obj)
 
 #==============================================================================
@@ -327,3 +333,43 @@ class DictItems(Iterable):
         """
         item = DictPopitem(self._dict_obj)
         return [IndexedElement(item, 0), IndexedElement(item, 1)]
+
+#==============================================================================
+class DictGetItem(DictMethod):
+    """
+    Represents a call to the .__getitem__() method.
+
+    The __getitem__() method returns the value for the specified key.
+
+    Parameters
+    ----------
+    dict_obj : TypedAstNode
+        The object from which the method is called.
+
+    k : TypedAstNode
+        The key which is used to select the value from the dictionary.
+    """
+    __slots__ = ('_class_type', '_shape')
+    name = 'get'
+
+    def __init__(self, dict_obj, k):
+        dict_type = dict_obj.class_type
+        self._class_type = dict_type.value_type
+        self._shape = None
+        if k.class_type != dict_type.key_type:
+            raise TypeError(f"Key passed to get method has type {k.class_type}. Expected {dict_type.key_type}")
+
+        super().__init__(dict_obj, k)
+
+        if self._class_type.rank:
+            self._shape = [PyccelArrayShapeElement(self,LiteralInteger(i)) \
+                    for i in range(self._class_type.rank)]
+
+    @property
+    def key(self):
+        """
+        The key that is used to select the element from the dict.
+
+        The key that is used to select the element from the dict.
+        """
+        return self._args[0]
diff --git a/pyccel/ast/class_defs.py b/pyccel/ast/class_defs.py
index 89a74ac0..386e6438 100644
--- a/pyccel/ast/class_defs.py
+++ b/pyccel/ast/class_defs.py
@@ -13,7 +13,7 @@ from pyccel.ast.builtin_methods.list_methods import (ListAppend, ListInsert, Lis
                                                      ListClear, ListExtend, ListRemove,
                                                      ListCopy, ListSort)
 from pyccel.ast.builtin_methods.dict_methods import (DictPop, DictPopitem, DictGet, DictClear,DictCopy,
-                                                     DictSetDefault, DictItems)
+                                                     DictSetDefault, DictItems, DictGetItem)
 
 from .builtins   import PythonImag, PythonReal, PythonConjugate
 from .core       import ClassDef, PyccelFunctionDef
@@ -187,6 +187,7 @@ DictClass = ClassDef('dict',
             PyccelFunctionDef('pop', func_class = DictPop),
             PyccelFunctionDef('popitem', func_class = DictPopitem),
             PyccelFunctionDef('setdefault', func_class = DictSetDefault),
+            PyccelFunctionDef('__getitem__', func_class = DictGetItem),
         ])
 
 #=======================================================================================
diff --git a/pyccel/codegen/printing/ccode.py b/pyccel/codegen/printing/ccode.py
index 23a111ec..964c1fd0 100644
--- a/pyccel/codegen/printing/ccode.py
+++ b/pyccel/codegen/printing/ccode.py
@@ -2865,6 +2865,21 @@ class CCodePrinter(CodePrinter):
 
         return pop_expr
 
+    def _print_DictGetItem(self, expr):
+        dict_obj = expr.dict_obj
+        dict_obj_code = self._print(ObjectAddress(dict_obj))
+        container_type = self.get_c_type(dict_obj.class_type)
+        key = self._print(expr.key)
+        assign = expr.get_user_nodes(Assign)
+        if assign:
+            assert len(assign) == 1
+            assign_node = assign[0]
+            lhs = assign_node.lhs
+            if lhs == expr or lhs.is_user_of(expr):
+                return f"(*{container_type}_at_mut({dict_obj_code}, {key}))"
+
+        return f"(*{container_type}_at({dict_obj_code}, {key}))"
+
     #=================== MACROS ==================
 
     def _print_MacroShape(self, expr):
diff --git a/pyccel/codegen/printing/fcode.py b/pyccel/codegen/printing/fcode.py
index 678f4483..90d733d2 100644
--- a/pyccel/codegen/printing/fcode.py
+++ b/pyccel/codegen/printing/fcode.py
@@ -1450,6 +1450,11 @@ class FCodePrinter(CodePrinter):
         var = self._print(expr.dict_obj)
         return f'call {var} % clear()\n'
 
+    def _print_DictGetItem(self, expr):
+        dict_obj = self._print(expr.dict_obj)
+        key = self._print(expr.key)
+        return f'{dict_obj} % of( {key} )'
+
     #========================== Numpy Elements ===============================#
 
     def _print_NumpySum(self, expr):
diff --git a/pyccel/codegen/printing/pycode.py b/pyccel/codegen/printing/pycode.py
index 19e22547..15c12f0c 100644
--- a/pyccel/codegen/printing/pycode.py
+++ b/pyccel/codegen/printing/pycode.py
@@ -872,15 +872,20 @@ class PythonCodePrinter(CodePrinter):
         key = self._print(expr.key)
         if expr.default_value:
             val = self._print(expr.default_value)
-            return f"{dict_obj}.get({key}, {val})\n"
+            return f"{dict_obj}.get({key}, {val})"
         else:
-            return f"{dict_obj}.get({key})\n"
+            return f"{dict_obj}.get({key})"
 
     def _print_DictItems(self, expr):
         dict_obj = self._print(expr.variable)
 
         return f"{dict_obj}.items()"
 
+    def _print_DictGetItem(self, expr):
+        dict_obj = self._print(expr.dict_obj)
+        key = self._print(expr.key)
+        return f"{dict_obj}[{key}]"
+
     def _print_Slice(self, expr):
         start = self._print(expr.start) if expr.start else ''
         stop  = self._print(expr.stop)  if expr.stop  else ''
diff --git a/pyccel/parser/semantic.py b/pyccel/parser/semantic.py
index b8be310e..3673676b 100644
--- a/pyccel/parser/semantic.py
+++ b/pyccel/parser/semantic.py
@@ -35,6 +35,7 @@ from pyccel.ast.builtins import Lambda, PythonMap
 
 from pyccel.ast.builtin_methods.list_methods import ListMethod, ListAppend
 from pyccel.ast.builtin_methods.set_methods  import SetAdd, SetUnion, SetCopy, SetIntersectionUpdate
+from pyccel.ast.builtin_methods.dict_methods  import DictGetItem, DictGet
 
 from pyccel.ast.core import Comment, CommentBlock, Pass
 from pyccel.ast.core import If, IfSection
@@ -692,6 +693,8 @@ class SemanticParser(BasicParser):
             self._indicate_pointer_target(pointer, target.lhs, expr)
         elif isinstance(target, IndexedElement):
             self._indicate_pointer_target(pointer, target.base, expr)
+        elif isinstance(target, (DictGetItem, DictGet)):
+            self._indicate_pointer_target(pointer, target.dict_obj, expr)
         elif isinstance(target, Variable):
             if target.is_alias:
                 try:
@@ -767,6 +770,11 @@ class SemanticParser(BasicParser):
             d_var['memory_handling'] = 'alias' if isinstance(var, Variable) else 'heap'
             return d_var
 
+        elif isinstance(expr, (DictGetItem, DictGet)):
+
+            d_var['memory_handling'] = 'alias' if not isinstance(expr.class_type, FixedSizeNumericType) else 'stack'
+            return d_var
+
         elif isinstance(expr, TypedAstNode):
 
             d_var['memory_handling'] = 'heap' if expr.rank > 0 else 'stack'
@@ -2740,28 +2748,41 @@ class SemanticParser(BasicParser):
         if isinstance(var, (PyccelFunctionDef, VariableTypeAnnotation, UnionTypeAnnotation)):
             return self._get_indexed_type(var, expr.indices, expr)
 
-        # TODO check consistency of indices with shape/rank
-        args = [self._visit(idx) for idx in expr.indices]
+        class_type = var.class_type
 
-        if (len(args) == 1 and isinstance(getattr(args[0], 'class_type', None), TupleType)):
-            args = args[0]
+        if isinstance(class_type, (NumpyNDArrayType, HomogeneousListType, TupleType)):
+            # TODO check consistency of indices with shape/rank
+            args = [self._visit(idx) for idx in expr.indices]
 
-        elif any(isinstance(getattr(a, 'class_type', None), TupleType) for a in args):
-            n_exprs = None
-            for a in args:
-                if getattr(a, 'shape', None) and isinstance(a.shape[0], LiteralInteger):
-                    a_len = a.shape[0]
-                    if n_exprs:
-                        assert n_exprs == a_len
-                    else:
-                        n_exprs = a_len
+            if (len(args) == 1 and isinstance(getattr(args[0], 'class_type', None), TupleType)):
+                args = args[0]
 
-            if n_exprs is not None:
-                new_expr_args = [[a[i] if hasattr(a, '__getitem__') else a for a in args]
-                                 for i in range(n_exprs)]
-                return NumpyArray(PythonTuple(*[var[a] for a in new_expr_args]))
+            elif any(isinstance(getattr(a, 'class_type', None), TupleType) for a in args):
+                n_exprs = None
+                for a in args:
+                    if getattr(a, 'shape', None) and isinstance(a.shape[0], LiteralInteger):
+                        a_len = a.shape[0]
+                        if n_exprs:
+                            assert n_exprs == a_len
+                        else:
+                            n_exprs = a_len
 
-        return self._extract_indexed_from_var(var, args, expr)
+                if n_exprs is not None:
+                    new_expr_args = [[a[i] if hasattr(a, '__getitem__') else a for a in args]
+                                     for i in range(n_exprs)]
+                    return NumpyArray(PythonTuple(*[var[a] for a in new_expr_args]))
+
+            return self._extract_indexed_from_var(var, args, expr)
+        else:
+            cls_base = self.scope.find(str(class_type), 'classes') or get_cls_base(class_type)
+            method = cls_base.get_method('__getitem__')
+            if method:
+                class_args = self._handle_function_args([FunctionCallArgument(a) for a in expr.indices])
+                args = [FunctionCallArgument(var), *class_args]
+                return self._handle_function(expr, method, args)
+            else:
+                raise errors.report(f"No __getitem__ found for type {class_type}",
+                        severity='fatal', symbol=expr)
 
     def _visit_PyccelSymbol(self, expr):
         name = expr
@@ -3524,10 +3545,12 @@ class SemanticParser(BasicParser):
 
         if isinstance(lhs, Variable):
             is_pointer = lhs.is_alias
-        elif isinstance(lhs, IndexedElement):
+        elif isinstance(lhs, (IndexedElement, DictGetItem)):
             is_pointer = False
         elif isinstance(lhs, (PythonTuple, PythonList)):
             is_pointer = any(l.is_alias for l in lhs if isinstance(lhs, Variable))
+        else:
+            raise NotImplementedError()
 
         # TODO: does is_pointer refer to any/all or last variable in list (currently last)
         is_pointer = is_pointer and isinstance(rhs, (Variable, Duplicate))


Test patch:
diff --git a/tests/epyccel/test_epyccel_dicts.py b/tests/epyccel/test_epyccel_dicts.py
index ffaa7f07..b29990db 100644
--- a/tests/epyccel/test_epyccel_dicts.py
+++ b/tests/epyccel/test_epyccel_dicts.py
@@ -193,6 +193,18 @@ def test_get_default_element(stc_language):
     assert isinstance(python_result, type(pyccel_result))
     assert python_result == pyccel_result
 
+def test_get_array(python_only_language):
+    def get_array():
+        import numpy as np
+        a = {1:np.ones(6), 2:np.zeros(4)}
+        tmp = a.get(1)
+        return tmp[1]
+    epyc_array = epyccel(get_array, language = python_only_language)
+    pyccel_result = epyc_array()
+    python_result = get_array()
+    assert isinstance(python_result, type(pyccel_result))
+    assert python_result == pyccel_result
+
 def test_get_str_keys(python_only_language):
     def get_str_keys():
         a = {'a':1, 'b':2}
@@ -233,6 +245,49 @@ def test_get_falsy_bool_default_element(stc_language):
     assert isinstance(python_result, type(pyccel_result))
     assert python_result == pyccel_result
 
+def test_getitem_element(language):
+    def getitem_element():
+        a = {1:1.0, 2:2.0}
+        return a[1]
+    epyc_element = epyccel(getitem_element, language = language)
+    pyccel_result = epyc_element()
+    python_result = getitem_element()
+    assert isinstance(python_result, type(pyccel_result))
+    assert python_result == pyccel_result
+
+def test_getitem_str_keys(python_only_language):
+    def getitem_str_keys():
+        a = {'a':1, 'b':2}
+        return a['a']
+    epyc_str_keys = epyccel(getitem_str_keys, language = python_only_language)
+    pyccel_result = epyc_str_keys()
+    python_result = getitem_str_keys()
+    assert isinstance(python_result, type(pyccel_result))
+    assert python_result == pyccel_result
+
+def test_getitem_array_element(python_only_language):
+    def getitem_array_element():
+        import numpy as np
+        a = {1:np.ones(6), 2:np.zeros(4)}
+        tmp = a[1]
+        return tmp[2]
+    epyc_array_element = epyccel(getitem_array_element, language = python_only_language)
+    pyccel_result = epyc_array_element()
+    python_result = getitem_array_element()
+    assert isinstance(python_result, type(pyccel_result))
+    assert python_result == pyccel_result
+
+def test_getitem_modify_element(language):
+    def getitem_modify_element():
+        a = {1:1.0, 2:2.0}
+        a[1] = 3.0
+        return a[1]
+    epyc_modify_element = epyccel(getitem_modify_element, language = language)
+    pyccel_result = epyc_modify_element()
+    python_result = getitem_modify_element()
+    assert isinstance(python_result, type(pyccel_result))
+    assert python_result == pyccel_result
+
 def test_dict_contains(language):
     def dict_contains():
         a = {1:1.0, 2:2.0, 3:3.0}