core.Source, FillCompute and FillSeq and core.Split support static co…

…ntext.

lena/core/fill_compute_seq.py

@@ -79,11 +79,15 @@ def __init__(self, *args):
         could not be correctly initialized,
         :exc:`.LenaTypeError` is raised.
         """
+        self._name = "FillComputeSeq"
+        super(FillComputeSeq, self).__init__(*args)
+        seq = self._seq
+
         before = []
         fc_el = None
         after = []
 
-        for ind, el in enumerate(args):
+        for ind, el in enumerate(seq):
             if not check_sequence_type.is_fill_compute_el(el):
                 before.append(el)
             else:
@@ -97,7 +101,7 @@ def __init__(self, *args):
             )
         self._fill_compute = fc_el
 
-        for el in args[ind+1:]:
+        for el in seq[ind+1:]:
             after.append(el)
 
         before.append(fc_el)
@@ -111,11 +115,6 @@ def __init__(self, *args):
         # or skip like here?
         self._after = sequence.Sequence(*after)
 
-        self._name = "FillComputeSeq"
-        seq = list(before_seq)
-        seq.extend(self._after._seq)
-        super(FillComputeSeq, self).__init__(*seq)
-
     def fill(self, value):
         """Fill *self* with *value*.
 

lena/core/fill_seq.py

@@ -41,19 +41,24 @@ def __init__(self, *args):
         or if an *args* is empty,
         :exc:`.LenaTypeError` is raised.
         """
-        seq = []
-
         if not args:
             raise exceptions.LenaTypeError(
                 "FillSeq must have at least one element"
             )
-        if not callable(getattr(args[-1], "fill", None)):
+
+        self._name = "FillSeq"  # for repr
+        super(FillSeq, self).__init__(*args)
+
+        seq = []
+        last = self._seq[-1]
+
+        if not callable(getattr(last, "fill", None)):
             raise exceptions.LenaTypeError(
                 "the last argument must implement fill method, "
-                "{} given".format(args[-1])
+                "{} given".format(last)
             )
         # convert all elements except last to FillInto (if needed)
-        for el in args[:-1]:
+        for el in self._seq[:-1]:
             if hasattr(el, "fill_into") and callable(el.fill_into):
                 seq.append(el)
             else:
@@ -68,18 +73,16 @@ def __init__(self, *args):
                     )
                 else:
                     seq.append(fill_into_el)
-        seq.append(args[-1])
+        seq.append(last)
         # transform FillInto elements into _Fill
-        fill_el = args[-1]
+        fill_el = last
         for el in reversed(seq[:-1]):
             fill_el = _Fill(el, fill_el)
         # note that self._seq consists of original FillInto elements.
         self._fill_el = fill_el
         # self for these methods is different
         self.fill = fill_el.fill
-
-        self._name = "FillSeq"  # for repr
-        super(FillSeq, self).__init__(*seq)
+        self._seq = seq
 
     def fill(self, value):
         """Fill *value* into an *element*.

lena/core/lena_sequence.py

@@ -12,9 +12,9 @@ class LenaSequence(object):
     get its length and get an item at the given index.
     """
     def __init__(self, *args):
+        from lena.context import update_recursively
         # todo: this doesn't require jinja2,
         # but we may want it to support template strings soon
-        from lena.context import update_recursively
         self._full_seq = args
         seq = []
 
@@ -88,10 +88,15 @@ def __repr__(self):
         return self._repr_nested()
 
     def _set_context(self, context):
+        from lena.context import update_recursively
+        # parent context doesn't necessarily has our context
+        # (for example, when we are inside Split)
+        update_recursively(self._context, context)
+        cont = self._context
         for el in self._need_context:
             # we don't use self._context,
             # because it was already set in this seq.
-            el._set_context(context)
+            el._set_context(cont)
 
     def _get_context(self):
         return deepcopy(self._context)

lena/core/source.py

@@ -34,23 +34,23 @@ def __init__(self, *args):
             raise exceptions.LenaTypeError(
                 "Source must be initialized with 1 argument or more (0 given)"
             )
-        if not callable(args[0]):
+
+        self._name = "Source"  # for repr
+        super(Source, self).__init__(*args)
+
+        first = self._seq[0]
+        if not callable(first):
             raise exceptions.LenaTypeError(
-                "first element {} ".format(args[0])
+                "first element {} ".format(first)
                 + "must be callable"
             )
-        self._first = args[0]
+        self._first = first
 
-        seq = [self._first]
         if len(args) > 1:
-            self._sequence = sequence.Sequence(*args[1:])
-            seq.extend(args[1:])
+            self._sequence = sequence.Sequence(*(self._seq[1:]))
         else:
             self._sequence = ()
 
-        self._name = "Source"  # for repr
-        super(Source, self).__init__(*seq)
-
     def __call__(self):
         """Generate flow."""
         flow = self._first()

lena/core/split.py

@@ -2,6 +2,7 @@
 import copy
 import itertools
 
+import lena.context
 from . import fill_compute_seq
 from . import check_sequence_type as ct
 from . import fill_request_seq
@@ -353,6 +354,22 @@ def repr_maybe_nested(el, base_indent, indent):
         return "".join([base_indent, "Split",
                         "([", mnl, elems, mnl, mbi, "])"])
 
+    def _get_context(self):
+        # we don't set context during initialisation,
+        # because we may need it at most once
+        # (during its parent sequence initialisation)
+        contexts = []
+        # the order of sequences is not important
+        for seq in self._seqs:
+            if hasattr(seq, "_get_context"):
+                contexts.append(seq._get_context())
+        return lena.context.intersection(*contexts)
+
+    def _set_context(self, context):
+        for seq in self._seqs:
+            if hasattr(seq, "_set_context"):
+                seq._set_context(context)
+
     def __eq__(self, other):
         if not isinstance(other, Split):
             return NotImplemented

tests/core/test_fill_seq.py

@@ -17,7 +17,7 @@ def test_lena_sequence_fill():
     assert len(s1) == 1
     assert len(s2) == 2
     # can get item
-    assert isinstance(s2[0], FillInto)
+    assert isinstance(s2._seq[0], FillInto)
     # deletion is prohibited
     with pytest.raises(TypeError):
         del s2[0]

tests/meta/test_elements.py

@@ -2,11 +2,22 @@
 
 import pytest
 
-from lena.core import Sequence, Source
+from lena.core import Sequence, Source, Split
 
 from lena.meta.elements import SetContext, UpdateContextFromStatic
 
 
+class StoreContext():
+
+    def __init__(self, name=""):
+        self._name = name
+        self._has_no_data = True
+
+    def _set_context(self, context):
+        print("StoreContext({}), {}".format(self._name, context))
+        self.context = context
+
+
 def test_set_context():
     data = [(0, {})]
     set_context_far = SetContext("data.detector", "far")
@@ -21,7 +32,7 @@ def test_set_context():
 
     # outside a sequence
     set_context_near = SetContext("data.detector", "near")
-    s1 = Sequence(set_context_near, s0)
+    s1 = Sequence(set_context_near, deepcopy(s0))
     assert list(s1.run(deepcopy(data))) == data
 
     # static context is updated after the element
@@ -40,3 +51,44 @@ def test_set_context():
     s1u0 = Sequence(s1, UpdateContextFromStatic())
     res1u0 = list(s1u0.run(deepcopy(data)))
     assert res1u0[0][1] == {'data': {'detector': 'near'}}
+
+
+def test_set_context_split():
+    set_context_far = SetContext("data.detector", "far")
+    set_context_near = SetContext("data.detector", "near")
+    set_context_common = SetContext("data.lost", True)
+    call = lambda _: "we won't run this"
+    store1, store2, store3, store4 = [StoreContext(str(i)) for i in range(1, 5)]
+
+    split = Split([
+        (
+            set_context_common,
+            call,
+            set_context_far,
+            store2,
+        ),
+        (
+            set_context_common,
+            call,
+            store3,
+            set_context_near,
+        ),
+    ])
+    assert split._get_context() == {'data': {'lost': True}}
+
+    s0 = Source(
+        SetContext("data.cycle", 1),
+        call,
+        store1,
+        split,
+        store4
+    )
+    assert store1.context == {'data': {'cycle': 1, 'lost': True}}
+    assert store2.context == {
+        'data': {'cycle': 1, 'detector': 'far', 'lost': True}
+    }
+    assert store3.context == {
+        'data': {'cycle': 1, 'detector': 'near', 'lost': True}
+    }
+    # static context is the same for the same level of nesting
+    assert store4.context == store1.context