Merge pull request #1350 from mathics/Cases-Heads-option

More Cases Heads option handling

Author: rocky

CHANGES.rst

@@ -27,6 +27,7 @@ Enhancements
 
 * a function `evaluate_predicate` allows for a basic predicate evaluation using `$Assumptions`.
 * ``Attributes`` accepts a string parameter.
+* ``Cases`` accepts Heads option. Issue #1302.
 * ``ColorNegate`` for colors is supported.
 * ``D`` and ``Derivative`` improvements.
 * ``FileNames`` returns a sorted list (#1250).

mathics/builtin/lists.py

@@ -25,6 +25,7 @@
     MessageException,
     NegativeIntegerException,
     CountableInteger,
+
 )
 from mathics.core.expression import (
     Expression,
@@ -50,7 +51,7 @@
 from mathics.core.expression import min_prec, machine_precision
 from mathics.core.expression import structure
 from mathics.core.evaluation import BreakInterrupt, ContinueInterrupt, ReturnInterrupt
-from mathics.core.rules import Pattern
+from mathics.core.rules import Pattern, Rule
 from mathics.core.convert import from_sympy
 from mathics.builtin.numbers.algebra import cancel
 from mathics.algorithm.introselect import introselect
@@ -1970,17 +1971,29 @@ def apply_pattern(self, items, sel, pattern, evaluation):
 class Cases(Builtin):
     """
     <dl>
-    <dt>'Cases[$list$, $pattern$]'
-        <dd>returns the elements of $list$ that match $pattern$.
-    <dt>'Cases[$list$, $pattern$, $ls$]'
-        <dd>returns the elements matching at levelspec $ls$.
+      <dt>'Cases[$list$, $pattern$]'
+      <dd>returns the elements of $list$ that match $pattern$.
+
+      <dt>'Cases[$list$, $pattern$, $ls$]'
+      <dd>returns the elements matching at levelspec $ls$.
+
+      <dt>'Cases[$list$, $pattern$, Head->$bool$]'
+      <dd>Match including the head of the expression in the search.
     </dl>
 
     >> Cases[{a, 1, 2.5, "string"}, _Integer|_Real]
      = {1, 2.5}
     >> Cases[_Complex][{1, 2I, 3, 4-I, 5}]
      = {2 I, 4 - I}
 
+    Find symbols among the elements of an expression:
+    >> Cases[{b, 6, \[Pi]}, _Symbol]
+     = {b, Pi}
+
+    Also include the head of the expression in the previous search:
+    >> Cases[{b, 6, \[Pi]}, _Symbol, Heads -> True]
+     = {List, b, Pi}
+
     #> Cases[1, 2]
      = {}
 
@@ -2008,24 +2021,33 @@ class Cases(Builtin):
         "Cases[pattern_][list_]": "Cases[list, pattern]",
     }
 
-    options = {"Heads": "False",}
+    options = {
+        "Heads": "False",
+    }
 
     def apply(self, items, pattern, ls, evaluation, options):
         "Cases[items_, pattern_, ls_:{1}, OptionsPattern[]]"
         if items.is_atom():
             return Expression(SymbolList)
 
+        from mathics.builtin.patterns import Matcher
+        if ls.has_form("Rule", 2):
+            if ls.leaves[0].get_name() == "System`Heads":
+                heads = ls.leaves[1].is_true()
+                ls = Expression("List", 1)
+            else:
+                return evaluation.message("Position", "level", ls)
+        else:
+            heads = self.get_option(options, "Heads", evaluation).is_true()
+
         try:
             start, stop = python_levelspec(ls)
         except InvalidLevelspecError:
             return evaluation.message("Position", "level", ls)
 
         results = []
 
-        from mathics.builtin.patterns import Matcher
-
         if pattern.has_form("Rule", 2) or pattern.has_form("RuleDelayed", 2):
-            from mathics.core.rules import Rule
 
             match = Matcher(pattern.leaves[0]).match
             rule = Rule(pattern.leaves[0], pattern.leaves[1])
@@ -2045,10 +2067,6 @@ def callback(level):
                     results.append(level)
                 return level
 
-        # TODO
-        heads = self.get_option(options, 'Heads', evaluation).is_true()
-        # heads = False
-
         walk_levels(items, start, stop, heads=heads, callback=callback)
 
         return Expression(SymbolList, *results)