Improve builtin Lookup

mathics/builtin/list/associations.py

@@ -8,16 +8,16 @@
 actual keys found in the collection.
 """
 
-
 from mathics.builtin.box.layout import RowBox
 from mathics.core.atoms import Integer
-from mathics.core.attributes import A_HOLD_ALL_COMPLETE, A_PROTECTED
+from mathics.core.attributes import A_HOLD_ALL_COMPLETE, A_PROTECTED, A_READ_PROTECTED
 from mathics.core.builtin import Builtin, Test
 from mathics.core.convert.expression import to_mathics_list
 from mathics.core.evaluation import Evaluation
 from mathics.core.expression import Expression
 from mathics.core.symbols import Symbol, SymbolTrue
 from mathics.core.systemsymbols import SymbolAssociation, SymbolMakeBoxes, SymbolMissing
+from mathics.eval.list.associations import eval_Lookup, eval_Lookup_multiple_keys
 from mathics.eval.lists import list_boxes
 
 
@@ -249,18 +249,60 @@ class Lookup(Builtin):
     <dl>
       <dt>Lookup[$assoc$, $key$]
       <dd>looks up the value associated with $key$ in the association $assoc$, \
-          or Missing[$KeyAbsent$].
+          returning Missing[$KeyAbsent$, $key$] if the key is not found.
+      <dt>Lookup[$assoc$, $key$, $default$]
+      <dd>looks up the value associated with $key$ in the association $assoc$, \
+          returning $default$ if the key is not found.
+      <dt>Lookup[$assoc$, {$key_1$, $key_2$, ...}]
+      <dd>looks up multiple keys and returns a list of values.
     </dl>
+
+    Look up the value associagted with key a:
+    >> Lookup[<|a -> 1, b -> 2|>, a]
+     = 1
+
+    When a key is not found, a Missing object is returned by default:
+    >> Lookup[<|a -> 1, b -> 2|>, c]
+     = Missing[KeyAbsent, c]
+
+    Provide a default value to be used when the key is not found:
+    >> Lookup[<|a -> 1, b -> 2|>, c, -1]
+     = -1
+
+    Use the operator form of Lookup:
+    >> Lookup[<|a -> 1, b -> 2|>, {a, b}]
+     = {1, 2}
+
+    Look up multiple keys at once:
+    >> Lookup[<|a -> 1, b -> 2|>, {a, b, c}]
+     = {1, 2, Missing[KeyAbsent, c]}
+
     """
 
-    attributes = A_HOLD_ALL_COMPLETE
-    rules = {
-        "Lookup[assoc_?AssociationQ, key_, default_]": "FirstCase[assoc, _[Verbatim[key], val_] :> val, default]",
-        "Lookup[assoc_?AssociationQ, key_]": 'Lookup[assoc, key, Missing["KeyAbsent", key]]',
+    attributes = A_PROTECTED | A_READ_PROTECTED
+
+    messages = {
+        "invrl": "The argument `1` is not a valid Association or a list of rules.",
     }
 
     summary_text = "perform lookup of a value by key, returning a specified default if it is not found"
 
+    def eval_assoc_key(self, assoc, key, evaluation: Evaluation):
+        """Lookup[assoc_Association, key_]"""
+        return eval_Lookup(assoc, key, None, evaluation)
+
+    def eval_assoc_key_default(self, assoc, key, default, evaluation: Evaluation):
+        """Lookup[assoc_Association, key_, default_]"""
+        return eval_Lookup(assoc, key, default, evaluation)
+
+    def eval_assoc_keys(self, assoc, keys, evaluation: Evaluation):
+        """Lookup[assoc_Association, keys_List]"""
+        return eval_Lookup_multiple_keys(assoc, keys, None, evaluation)
+
+    def eval_assoc_keys_default(self, assoc, keys, default, evaluation: Evaluation):
+        """Lookup[assoc_Association, keys_List, default_]"""
+        return eval_Lookup_multiple_keys(assoc, keys, default, evaluation)
+
 
 class Missing(Builtin):
     """

mathics/core/systemsymbols.py

@@ -184,6 +184,7 @@
 SymbolIntegrate = Symbol("System`Integrate")
 SymbolInterpretationBox = Symbol("System`InterpretationBox")
 SymbolKey = Symbol("System`Key")
+SymbolKeyAbsent = Symbol("System`KeyAbsent")
 SymbolKhinchin = Symbol("System`Khinchin")
 SymbolLabelStyle = Symbol("System`LabelStyle")
 SymbolLast = Symbol("System`Last")

mathics/eval/list/__init__.py

@@ -0,0 +1,4 @@
+"""
+Evaluation routines and associated code for Built-in functions found under module
+mathics.builtins.list.
+"""

mathics/eval/list/associations.py

@@ -0,0 +1,50 @@
+from mathics.core.evaluation import Evaluation
+from mathics.core.expression import Expression
+from mathics.core.list import ListExpression
+from mathics.core.systemsymbols import SymbolKeyAbsent, SymbolMissing
+
+
+def eval_Lookup(assoc, key, default, evaluation: Evaluation):
+    """Evaluation method for Lookup."""
+
+    if assoc.has_form("Association", None):
+        # Search through association elements (rules)
+        for element in assoc.elements:
+            if element.has_form(("Rule", "RuleDelayed"), 2):
+                if element.elements[0] == key:
+                    return element.elements[1]
+
+        # Key not found
+        if default is not None:
+            return default
+        else:
+            return Expression(SymbolMissing, SymbolKeyAbsent, key)
+
+    elif isinstance(assoc, ListExpression):
+        # Search through list of rules
+        for element in assoc.elements:
+            if element.has_form(("Rule", "RuleDelayed"), 2):
+                if element.elements[0] == key:
+                    return element.elements[1]
+
+        # Key not found
+        if default is not None:
+            return default
+        else:
+            return Expression(SymbolMissing, SymbolKeyAbsent, key)
+
+    elif assoc.has_form(("Rule", "RuleDelayed"), 2):
+        if assoc.elements[0] == key:
+            return assoc.elements[1]
+        return None
+
+    else:
+        evaluation.message("Lookup", "invrl", assoc)
+        # Should we return SymbolFailed?
+        return None
+
+
+def eval_Lookup_multiple_keys(assoc, keys, default, evaluation: Evaluation):
+    """Evaluation method for Lookup with multiple keys, threading over the key list."""
+    results = [eval_Lookup(assoc, key, default, evaluation) for key in keys.elements]
+    return ListExpression(*results)

test/builtin/list/test_association.py

@@ -263,3 +263,24 @@ def test_map_over_associations(
         failure_message=assert_message,
         expected_messages=expected_messages,
     )
+
+
+@pytest.mark.parametrize(
+    ("str_expr", "expected_messages", "str_expected", "assert_message"),
+    [
+        (
+            'a=Association[{"F":>1,"G":>2}]; Lookup[a, "H"]',
+            None,
+            "Missing[KeyAbsent, H]",
+            "Lookup test on an association variable where the key is not found.",
+        ),
+        ("ClearAll[a];", None, "Null", None),
+    ],
+)
+def test_lookup(str_expr, expected_messages, str_expected, assert_message):
+    check_evaluation(
+        str_expr,
+        str_expected,
+        failure_message=assert_message,
+        expected_messages=expected_messages,
+    )

test/builtin/list/test_eol.py

@@ -10,6 +10,7 @@
 @pytest.mark.parametrize(
     ("str_expr", "expected_messages", "str_expected", "assert_message"),
     [
+        ("ClearAll[a];", None, "Null", None),
         (
             "Append[a, b]",
             ("Nonatomic expression expected at position 1 in Append[a, b].",),