Merge pull request #2645 from stfc/2588_get_routine_psyir

(Closes #2585) Add an optional, 'default' argument to SymbolTable.lookup()

changelog

@@ -169,6 +169,9 @@
 	59) PR #2649 towards #2586. Adds support for pointer assignments of
 	the form `var =[>] this[%...]%procedure(...)`.
 
+	60) PR #2645 for #2585. Add an optional "otherwise" argument to
+	the SymbolTalbe lookup method.
+
 release 2.5.0 14th of February 2024
 
 	1) PR #2199 for #2189. Fix bugs with missing maps in enter data

src/psyclone/psyir/symbols/symbol_table.py

@@ -56,6 +56,12 @@
 from psyclone.psyir.symbols.typed_symbol import TypedSymbol
 
 
+# Used to provide a unique default value for methods within the
+# SymbolTable class. This enables us to determine when the user has
+# supplied a value of 'None' for arguments that have a default value.
+DEFAULT_SENTINEL = object()
+
+
 class SymbolTable():
     # pylint: disable=too-many-public-methods
     '''Encapsulates the symbol table and provides methods to add new
@@ -933,32 +939,37 @@ def specify_argument_list(self, argument_symbols):
         self._validate_arg_list(argument_symbols)
         self._argument_list = argument_symbols[:]
 
-    def lookup(self, name, visibility=None, scope_limit=None):
+    def lookup(self, name, visibility=None, scope_limit=None,
+               otherwise=DEFAULT_SENTINEL):
         '''Look up a symbol in the symbol table. The lookup can be limited
         by visibility (e.g. just show public methods) or by scope_limit (e.g.
         just show symbols up to a certain scope).
 
         :param str name: name of the symbol.
-        :param visibilty: the visibility or list of visibilities that the \
+        :param visibilty: the visibility or list of visibilities that the
                           symbol must have.
         :type visibility: [list of] :py:class:`psyclone.symbols.Visibility`
-        :param scope_limit: optional Node which limits the symbol \
-            search space to the symbol tables of the nodes within the \
-            given scope. If it is None (the default), the whole \
-            scope (all symbol tables in ancestor nodes) is searched \
-            otherwise ancestors of the scope_limit node are not \
+        :param scope_limit: optional Node which limits the symbol
+            search space to the symbol tables of the nodes within the
+            given scope. If it is None (the default), the whole
+            scope (all symbol tables in ancestor nodes) is searched
+            otherwise ancestors of the scope_limit node are not
             searched.
-        :type scope_limit: :py:class:`psyclone.psyir.nodes.Node` or \
-            `NoneType`
+        :type scope_limit: Optional[:py:class:`psyclone.psyir.nodes.Node`]
+        :param Any otherwise: an optional value to return if the named symbol
+            cannot be found (rather than raising a KeyError).
 
         :returns: the symbol with the given name and, if specified, visibility.
+                  If no match is found and 'otherwise' is supplied then that
+                  value is returned.
         :rtype: :py:class:`psyclone.psyir.symbols.Symbol`
 
         :raises TypeError: if the name argument is not a string.
-        :raises SymbolError: if the name exists in the Symbol Table but does \
+        :raises SymbolError: if the name exists in the Symbol Table but does
                              not have the specified visibility.
         :raises TypeError: if the visibility argument has the wrong type.
-        :raises KeyError: if the given name is not in the Symbol Table.
+        :raises KeyError: if the given name is not in the Symbol Table and
+                          `otherwise` is not supplied.
 
         '''
         if not isinstance(name, str):
@@ -992,8 +1003,11 @@ def lookup(self, name, visibility=None, scope_limit=None):
                         f" match with the requested visibility: {vis_names}")
             return symbol
         except KeyError as err:
-            raise KeyError(f"Could not find '{name}' in the Symbol Table.") \
-                from err
+            if otherwise is DEFAULT_SENTINEL:
+                # No 'otherwise' value supplied so we raise an exception.
+                raise KeyError(f"Could not find '{name}' in the Symbol "
+                               f"Table.") from err
+            return otherwise
 
     def lookup_with_tag(self, tag, scope_limit=None):
         '''Look up a symbol by its tag. The lookup can be limited by

src/psyclone/tests/psyir/symbols/symbol_table_test.py

@@ -1230,7 +1230,7 @@ def test_lookup_1():
         str(error.value)
 
 
-def test_lookup_2():
+def test_lookup_vis_filtering():
     '''Test the visibility argument filtering functionality of the
     lookup() method.
 
@@ -1271,7 +1271,7 @@ def test_lookup_2():
             "searching for symbol 'var2'" in str(err.value))
 
 
-def test_lookup_3():
+def test_lookup_name_type_error():
     '''Check that lookup() in the SymbolTable class raises the expected
     exception if the name argument has the wrong type.
 
@@ -1284,7 +1284,7 @@ def test_lookup_3():
             "but found 'DataSymbol'." in str(info.value))
 
 
-def test_lookup_4():
+def test_lookup_scope_limit():
     '''Check that lookup() in the SymbolTable class behaves as
     expected with the scope_limit argument.
 
@@ -1315,6 +1315,28 @@ def test_lookup_4():
                 in str(info.value))
 
 
+def test_lookup_default_value():
+    '''Test that the lookup() method returns the supplied (optional)
+    'otherwise' value if a Symbol cannot be found.
+
+    '''
+    table = symbols.SymbolTable()
+    with pytest.raises(KeyError) as err:
+        table.lookup("missing")
+    assert "Could not find 'missing'" in str(err.value)
+    # The most common use case is to allow None to be returned rather than
+    # raising an error.
+    sym = table.lookup("missing", otherwise=None)
+    assert sym is None
+    # But any other alternative value can be specified.
+    sym = table.lookup("missing", otherwise=symbols.Symbol("missing"))
+    assert sym.name == "missing"
+    # The presence of other, optional arguments doesn't alter anything.
+    sym = table.lookup("missing", visibility=symbols.Symbol.Visibility.PUBLIC,
+                       otherwise=1)
+    assert sym == 1
+
+
 def test_lookup_with_tag_1():
     '''Test that the lookup_with_tag method retrieves symbols from the symbol
     table if the tag exists, otherwise it raises an error.'''