WriteROOTTree uses context.output instead of context.data. Add exampl…

…e for ReadROOTFile and ReadROOTTree.

docs/examples/topics/root/read_staff_tree.py

@@ -0,0 +1,77 @@
+import copy
+import itertools
+import sys
+
+import ROOT
+from ROOT import TFile, TTree, gROOT
+
+import lena
+from lena.flow import Print
+from lena.input import ReadROOTFile, ReadROOTTree
+from lena.variables import Variable
+
+
+def staff_get_entries():
+    # create a C++ structure
+    gROOT.ProcessLine(
+    "struct staff_t {\
+       Int_t           Category;\
+       UInt_t          Flag;\
+       Int_t           Age;\
+       Int_t           Service;\
+       Int_t           Children;\
+       Int_t           Grade;\
+       Int_t           Step;\
+       Int_t           Hrweek;\
+       Int_t           Cost;\
+       Char_t          Division[4];\
+       Char_t          Nation[3];\
+    };")
+    # get a staff_t object into Python
+    staff = ROOT.staff_t()
+    def get_entries(tree):
+        # here you can disable some branches to optimise performance
+        tree.SetBranchAddress("staff", staff)
+        for entry in tree:
+            # or copy.copy if there is no deep nesting
+            yield copy.deepcopy(staff)
+    return get_entries
+
+
+# imperative approach
+f = TFile("staff.root")
+staff_tree = f.Get("T")
+for staff_m in itertools.islice(staff_tree, 10):
+    print(staff_m.Age, end='')
+print()
+
+get_entries = staff_get_entries()
+for staff_m in itertools.islice(get_entries(staff_tree), 10):
+    print(staff_m.Age, ", ", sep="", end="")
+print()
+
+# Lena sequence
+s1 = lena.core.Sequence(
+    ReadROOTFile(),
+    ReadROOTTree(get_entries=staff_get_entries()),
+    Variable("age", lambda staff: staff.Age),
+    # create histograms, do normal Lena processing
+    Print(),
+    lena.flow.Slice(10),
+)
+
+# run the sequence
+list(s1.run(["staff.root"]))
+# (58, {'input': {'root_file_path': 'staff.root', 'root_tree_name': 'T'}, 'variable': {'name': 'age'}})
+# ...
+
+s2 = lena.core.Sequence(
+    ReadROOTFile(),
+    # several leaves can be read into a named tuple
+    ReadROOTTree(["Divisions/Division", "Age"]),
+    Print(),
+    lena.flow.Slice(10),
+)
+
+# run the sequence
+list(s2.run(["staff.root"]))

docs/source/output.rst

@@ -9,6 +9,7 @@ Output
     PDFToPNG
     ToCSV
     Write
+..    WriteROOTTree
 
 .. not covered yet: JSONEncoder, ExtendedHistToHist
 
@@ -33,6 +34,7 @@ Output
 .. autofunction:: hist2d_to_csv
 .. autoclass:: Write
 .. autoclass:: Writer
+.. .. autoclass:: WriteROOTTree
 
 LaTeX
 -----

docs/source/variables.rst

@@ -1,3 +1,5 @@
+.. _module_variables:
+
 Variables
 =========
 **Variables:**

lena/input/read_root_tree.py

@@ -41,6 +41,8 @@ def __init__(self, leaves=None, get_entries=None):
         # still enables "from lena.flow import ReadROOTTree",
         # instead of "from lena.flow.read_root_tree import ReadROOTTree"
         import ROOT
+        # todo: add tuple_name to kwargs
+        # (otherwise T_entry can look weird/frightening)
 
         if leaves is not None:
             err_msg = ""

lena/output/__init__.py

@@ -3,6 +3,7 @@
 from .make_filename import MakeFilename
 from .write import Writer, Write
 from .to_csv import ToCSV, hist1d_to_csv, hist2d_to_csv
+from .write_root_tree import WriteROOTTree
 
 
 def raise_on_usage(clsname, modname):
@@ -28,6 +29,7 @@ def stub(*args, **kwargs):
     'RenderLaTeX',
     'Write',
     'Writer',
+    'WriteROOTTree',
     'ToCSV', 'hist1d_to_csv', 'hist2d_to_csv',
     'RenderLaTeX'
 ]

lena/output/write_root_tree.py

@@ -8,12 +8,12 @@
 
 
 class WriteROOTTree():
-    """Write data from the incoming flow to a ROOT tree."""
+    """Write data to a ROOT tree."""
 
     def __init__(self, name, root_file):
         """*name* is the name of the tree.
 
-        Only tuples produced by lena :module:`Variables`
+        Only tuples produced by lena :ref:`module_variables`
         or namedtuples can be written to tree.
         They can contain only int and float Python types.
         For more general options, use a ROOT tree directly.
@@ -33,6 +33,10 @@ def __init__(self, name, root_file):
         to the same file several times.
         *Recreate* (default) means that file
         will be completely overwritten.
+
+        Note
+        ----
+            To use this class, `ROOT <https://root.cern/>`__ must be installed.
         """
         import ROOT
         # todo: name -> tree, so that a tree could be passed?
@@ -252,7 +256,7 @@ def run(self, flow):
         If an opened ROOT file was provided during the initialization,
         it is the user's responsibility to close that.
 
-        *Context.data* is updated with *root_file_path*
+        *Context.output* is updated with *root_file_path*
         and *root_tree_name*.
         """
         # otherwise the function won't know its name
@@ -300,7 +304,7 @@ def run(self, flow):
                 root_file.Close()
 
         lena.context.update_recursively(
-            context, {"data": copy.deepcopy(self._data_c)}
+            context, {"output": copy.deepcopy(self._data_c)}
         )
 
         yield (self._root_file_name, context)

tests/root/test_write_root_tree.py

@@ -58,7 +58,7 @@ def test_run():
     assert list(seq.run(data)) == [
         ('file.root',
             {
-                "data": {
+                "output": {
                     'root_file_path': 'file.root',
                     'root_tree_name': 'tree'
                 },