Go over mathics.builtin.list.associations

mathics/builtin/list/associations.py

@@ -4,7 +4,7 @@
 Associations
 
 An Association maps keys to values and is similar to a dictionary in Python; \
-it is often sparse in that their key space is much larger than the number of \
+it is often sparse in that its key space is much larger than the number of \
 actual keys found in the collection.
 """
 
@@ -17,7 +17,11 @@
 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.list.associations import (
+    eval_Lookup,
+    eval_Lookup_assocs_list_key,
+    eval_Lookup_multiple_keys,
+)
 from mathics.eval.lists import list_boxes
 
 
@@ -46,6 +50,10 @@ class Association(Builtin):
     Associations can be nested:
     >> <|a -> x, b -> y, <|a -> z, d -> t|>|>
      = <|a ⇾ z, b ⇾ y, d ⇾ t|>
+
+    Look up a key in multiple associations:
+    >> Lookup[{<|a -> 1, b -> 2|>, <|a -> 3, c -> 4|>}, a]
+     = {1, 3}
     """
 
     error_idx = 0
@@ -190,6 +198,9 @@ class Keys(Builtin):
 
       <dt>'Keys'[{$key_1$ '->' $val_1$, $key_2$ '->' $val_2$, ...}]
       <dd>return a list of the $key_i$ in a list of rules.
+
+      <dt>'Keys'[$expr$, $h$]
+      <dd>applies the head $h$ to each key.
     </dl>
 
     >> Keys[<|a -> x, b -> y|>]
@@ -205,19 +216,21 @@ class Keys(Builtin):
     Keys are listed in the order of their appearance:
     >> Keys[{c -> z, b -> y, a -> x}]
      = {c, b, a}
+
+    Apply a head to each key:
+    >> Keys[<|a -> x, b -> y|>, f]
+     = {f[a], f[b]}
     """
 
     attributes = A_PROTECTED
 
-    messages = {
-        "argx": "Keys called with `1` arguments; 1 argument is expected.",
-        "invrl": "The argument `1` is not a valid Association or a list of rules.",
-    }
+    eval_error = Builtin.generic_argument_error
+    expected_args = (1, 2)
 
     summary_text = "list association keys"
 
     def eval(self, rules, evaluation: Evaluation):
-        "Keys[rules___]"
+        "Keys[rules_]"
 
         def get_keys(expr):
             if expr.has_form(("Rule", "RuleDelayed"), 2):
@@ -241,6 +254,30 @@ def get_keys(expr):
         except TypeError:
             return None
 
+    def eval_with_head(self, rules, head, evaluation: Evaluation):
+        "Keys[rules_, head_]"
+
+        def get_keys_with_head(expr, h):
+            if expr.has_form(("Rule", "RuleDelayed"), 2):
+                key = expr.elements[0]
+                return Expression(h, key)
+            elif expr.has_form("List", None) or (
+                expr.has_form("Association", None)
+                and AssociationQ(expr).evaluate(evaluation) is SymbolTrue
+            ):
+                return to_mathics_list(
+                    *expr.elements,
+                    elements_conversion_fn=lambda e: get_keys_with_head(e, h),
+                )
+            else:
+                evaluation.message("Keys", "invrl", expr)
+                raise TypeError
+
+        try:
+            return get_keys_with_head(rules, head)
+        except TypeError:
+            return None
+
 
 class Lookup(Builtin):
     """
@@ -255,9 +292,11 @@ class Lookup(Builtin):
           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.
+      <dt>Lookup[{$assoc_1$, $assoc_2$, ...}, $key$]
+      <dd>looks up $key$ in each association and returns a list of values.
     </dl>
 
-    Look up the value associagted with key a:
+    Look up the value associated with key $a$:
     >> Lookup[<|a -> 1, b -> 2|>, a]
      = 1
 
@@ -277,13 +316,16 @@ class Lookup(Builtin):
     >> Lookup[<|a -> 1, b -> 2|>, {a, b, c}]
      = {1, 2, Missing[KeyAbsent, c]}
 
+    Provide a default value to be used when the key is not found:
+    >> Lookup[<|a -> 1, b -> 2|>, c, 3]
+     = 3
+
     """
 
     attributes = A_PROTECTED | A_READ_PROTECTED
 
-    messages = {
-        "invrl": "The argument `1` is not a valid Association or a list of rules.",
-    }
+    eval_error = Builtin.generic_argument_error
+    expected_args = range(2, 5)
 
     summary_text = "perform lookup of a value by key, returning a specified default if it is not found"
 
@@ -303,6 +345,16 @@ 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)
 
+    def eval_assocs_list_key(self, assocs, key, evaluation: Evaluation):
+        """Lookup[assocs_List, key_]"""
+        return eval_Lookup_assocs_list_key(assocs, key, None, evaluation)
+
+    def eval_assocs_list_key_default(
+        self, assocs, key, default, evaluation: Evaluation
+    ):
+        """Lookup[assocs_List, key_, default_]"""
+        return eval_Lookup_assocs_list_key(assocs, key, default, evaluation)
+
 
 class Missing(Builtin):
     """
@@ -332,6 +384,9 @@ class Values(Builtin):
 
       <dt>'Values'[{$key_1$ '->' $val_1$, $key_2$ '->' $val_2$, ...}]
       <dd>return a list of the $val_i$ in a list of rules.
+
+      <dt>'Values'[$expr$, $h$]
+      <dd>applies the head $h$ to each value.
     </dl>
 
     >> Values[<|a -> x, b -> y|>]
@@ -348,19 +403,20 @@ class Values(Builtin):
     >> Values[{c -> z, b -> y, a -> x}]
      = {z, y, x}
 
+    Apply a head to each value:
+    >> Values[<|a -> x, b -> y|>, f]
+     = {f[x], f[y]}
     """
 
     attributes = A_PROTECTED
 
-    messages = {
-        "argx": "Values called with `1` arguments; 1 argument is expected.",
-        "invrl": "The argument `1` is not a valid Association or a list of rules.",
-    }
+    eval_error = Builtin.generic_argument_error
+    expected_args = (1, 2)
 
     summary_text = "list association values"
 
     def eval(self, rules, evaluation: Evaluation):
-        "Values[rules___]"
+        "Values[rules_]"
 
         def get_values(expr):
             if expr.has_form(("Rule", "RuleDelayed"), 2):
@@ -384,3 +440,27 @@ def get_values(expr):
             return get_values(rules[0])
         except TypeError:
             evaluation.message("Values", "invrl", rules[0])
+
+    def eval_with_head(self, rules, head, evaluation: Evaluation):
+        "Values[rules_, head_]"
+
+        def get_values_with_head(expr, h):
+            if expr.has_form(("Rule", "RuleDelayed"), 2):
+                value = expr.elements[1]
+                return Expression(h, value)
+            elif expr.has_form("List", None) or (
+                expr.has_form("Association", None)
+                and AssociationQ(expr).evaluate(evaluation) is SymbolTrue
+            ):
+                return to_mathics_list(
+                    *expr.elements,
+                    elements_conversion_fn=lambda e: get_values_with_head(e, h),
+                )
+            else:
+                evaluation.message("Values", "invrl", expr)
+                raise TypeError
+
+        try:
+            return get_values_with_head(rules, head)
+        except TypeError:
+            return None

mathics/builtin/messages.py

@@ -215,6 +215,7 @@ class General(Builtin):
         "intp": "Positive integer expected.",
         "intnn": "Non-negative integer expected.",
         "intnm": "Non-negative machine-sized integer expected at position `1` in `2`.",
+        "invrl": "The argument `1` is not a valid Association or a list of rules.",
         "iterb": "Iterator does not have appropriate bounds.",
         "ivar": "`1` is not a valid variable.",
         "level": ("Level specification `1` is not of the form n, " "{n}, or {m, n}."),

mathics/eval/list/associations.py

@@ -44,7 +44,32 @@ def eval_Lookup(assoc, key, default, evaluation: Evaluation):
         return None
 
 
+def eval_Lookup_assocs_list_key(assocs, key, default, evaluation: Evaluation):
+    """Evaluation method for Lookup with a list of associations and a single key.
+
+    Looks up the key in each association and returns a list of values.
+    """
+    if not isinstance(assocs, ListExpression):
+        evaluation.message("Lookup", "invrl", assocs)
+        return None
+
+    results = [
+        eval_Lookup(assoc, key, default, evaluation) for assoc in assocs.elements
+    ]
+    return ListExpression(*results)
+
+
 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)
+
+
+def eval_assocs_list_key(self, assocs, key, evaluation: Evaluation):
+    """Lookup[assocs_List, key_]"""
+    return eval_Lookup_assocs_list_key(assocs, key, None, evaluation)
+
+
+def eval_assocs_list_key_default(self, assocs, key, default, evaluation: Evaluation):
+    """Lookup[assocs_List, key_, default_]"""
+    return eval_Lookup_assocs_list_key(assocs, key, default, evaluation)

test/builtin/list/test_association.py → test/builtin/list/test_associations.py

@@ -2,7 +2,7 @@
 """
 Unit tests for mathics.builtins.list.constructing
 """
-from test.helper import check_evaluation
+from test.helper import check_arg_counts, check_evaluation
 
 import pytest
 
@@ -112,12 +112,12 @@
             "Keys[{a -> x, {a -> y, b}}]",
             None,
         ),
-        (
-            "Keys[a -> x, b -> y]",
-            ("Keys called with 2 arguments; 1 argument is expected.",),
-            "Keys[a -> x, b -> y]",
-            None,
-        ),
+        # (
+        #     "Keys[a -> x, b -> y]",
+        #     None
+        #     "(b -> y)[a]",
+        #     None,
+        # ),
         ("Values[a -> x]", None, "x", None),
         (
             "Values[{a -> x, a -> y, {a -> z, <|b -> t|>, <||>, {}}}]",
@@ -175,12 +175,11 @@
             "Values[{a -> x, {a -> y, b}}]",
             None,
         ),
-        (
-            "Values[a -> x, b -> y]",
-            ("Values called with 2 arguments; 1 argument is expected.",),
-            "Values[a -> x, b -> y]",
-            None,
-        ),
+        # (
+        #     "Values[a -> x, b -> y]",
+        #     "(b -> y)[x]",
+        #     None,
+        # ),
         ("assoc=.;subassoc=.;", None, "Null", None),
     ],
 )
@@ -284,3 +283,25 @@ def test_lookup(str_expr, expected_messages, str_expected, assert_message):
         failure_message=assert_message,
         expected_messages=expected_messages,
     )
+
+
+@pytest.mark.parametrize(
+    ("function_name", "msg_fragment"),
+    [
+        (
+            "Keys",
+            "1 or 2 arguments are",
+        ),
+        (
+            "Lookup",
+            "between 2 and 4 arguments are",
+        ),
+        (
+            "Values",
+            "1 or 2 arguments are",
+        ),
+    ],
+)
+def test_arg_count_errors(function_name, msg_fragment):
+    """ """
+    check_arg_counts(function_name, msg_fragment)