Merge "generateMany function was added"

yaql/language/utils.py

@@ -57,6 +57,7 @@ MappingType = collections.Mapping
 MutableMappingType = collections.MutableMapping
 IterableType = collections.Iterable
 IteratorType = collections.Iterator
+QueueType = collections.deque
 
 
 def convert_input_data(obj, rec=None):

yaql/standard_library/queries.py

@@ -595,15 +595,53 @@ def accumulate(collection, selector, seed=utils.NO_VALUE):
 
 
 @specs.parameter('predicate', yaqltypes.Lambda())
-@specs.parameter('next_', yaqltypes.Lambda())
+@specs.parameter('producer', yaqltypes.Lambda())
 @specs.parameter('selector', yaqltypes.Lambda())
-def generate(initial, predicate, next_, selector=None):
+@specs.parameter('decycle', bool)
+def generate(engine, initial, predicate, producer, selector=None,
+             decycle=False):
+    past_items = None if not decycle else set()
     while predicate(initial):
+        if past_items is not None:
+            if initial in past_items:
+                break
+            past_items.add(initial)
+            utils.limit_memory_usage(engine, (1, past_items))
         if selector is None:
             yield initial
         else:
             yield selector(initial)
-        initial = next_(initial)
+        initial = producer(initial)
+
+
+@specs.parameter('producer', yaqltypes.Lambda())
+@specs.parameter('selector', yaqltypes.Lambda())
+@specs.parameter('decycle', bool)
+@specs.parameter('depth_first', bool)
+def generate_many(engine, initial, producer, selector=None, decycle=False,
+                  depth_first=False):
+    past_items = None if not decycle else set()
+    queue = utils.QueueType([initial])
+    while queue:
+        item = queue.popleft()
+        if past_items is not None:
+            if item in past_items:
+                continue
+            else:
+                past_items.add(item)
+                utils.limit_memory_usage(engine, (1, past_items))
+        if selector is None:
+            yield item
+        else:
+            yield selector(item)
+        produced = producer(item)
+        if depth_first:
+            len_before = len(queue)
+            queue.extend(produced)
+            queue.rotate(len(queue) - len_before)
+        else:
+            queue.extend(produced)
+        utils.limit_memory_usage(engine, (1, queue))
 
 
 @specs.method
@@ -676,4 +714,5 @@ def register(context):
     context.register_function(is_iterable)
     context.register_function(sequence)
     context.register_function(generate)
+    context.register_function(generate_many)
     context.register_function(default_if_empty)

yaql/tests/__init__.py

@@ -88,7 +88,9 @@ class TestCase(testtools.TestCase):
 
     def eval(self, expression, data=None, context=None):
         expr = self.engine(expression)
-        return expr.evaluate(data=data, context=context or self.context)
+        context = context or self.context
+        context['data'] = data
+        return expr.evaluate(data=data, context=context)
 
     def legacy_eval(self, expression, data=None, context=None):
         expr = self.legacy_engine(expression)

yaql/tests/test_queries.py

@@ -370,6 +370,35 @@ class TestQueries(yaql.tests.TestCase):
             [0, 4, 16, 36, 64],
             self.eval('generate(0, $ < 10, $ + 2, $ * $)'))
 
+    def test_generate_many(self):
+        friends = {
+            'John': ['Jim'],
+            'Jim': ['Jay', 'Jax'],
+            'Jax': ['John', 'Jacob', 'Jonathan'],
+            'Jacob': ['Jonathan', 'Jenifer'],
+        }
+        self.assertEqual(
+            ['John', 'Jim', 'Jay', 'Jax', 'Jacob', 'Jonathan', 'Jenifer'],
+            self.eval(
+                'generateMany(John, $data.get($, []), decycle => true)',
+                friends))
+
+        self.assertEqual(
+            ['John', 'Jim', 'Jay', 'Jax', 'Jacob', 'Jonathan', 'Jenifer'],
+            self.eval(
+                'generateMany(John, $data.get($, []), '
+                'decycle => true, depthFirst => true)', friends))
+
+        self.assertEqual(
+            ['Jay'],
+            self.eval('generateMany(Jay, $data.get($, []))', friends))
+
+        self.assertEqual(
+            ['JAX', 'JOHN', 'JACOB', 'JONATHAN', 'JIM', 'JENIFER', 'JAY'],
+            self.eval(
+                'generateMany(Jax, $data.get($, []), $.toUpper(), '
+                'decycle => true)', friends))
+
     def test_max(self):
         self.assertEqual(
             0,