Repository: rootpy/rootpy

Issue:
TreeChain draw fails: 'Graph' object has no attribute 'SetDirectory' 
When executing a 2D draw command (`chain.Draw("a:b")`) with TreeChain it fails with the following message:

```
Traceback (most recent call last):
  File "angular.py", line 8, in <module>
    tData.Draw("uE:dE", "")
  File "/home/munk/.local/lib/python2.7/site-packages/rootpy-0.8.0.dev0-py2.7.egg/rootpy/tree/chain.py", line 119, in Draw
    output.SetDirectory(0)
AttributeError: 'Graph' object has no attribute 'SetDirectory'
```

The caveat is, that unlike histograms, graphs is not associated with the current directory.


Discussion:
For the moment I'd suggest sticking with ROOT.TChain for Draw()-ing. rootpy's TreeChain is mainly beneficial when looping entry-by-entry. TreeChain's Draw can only handle histograms currently.

Resolving patch:
diff --git a/rootpy/plotting/graph.py b/rootpy/plotting/graph.py
--- a/rootpy/plotting/graph.py
+++ b/rootpy/plotting/graph.py
@@ -256,6 +256,7 @@ def __setitem__(self, index, point):
             raise IndexError("graph point index out of range")
         self.SetPoint(index, *point)
 
+
 class _Graph1DBase(_GraphBase):
 
     @classmethod
@@ -618,6 +619,26 @@ def Integrate(self):
             area += (X[i + 1] - X[i]) * (Y[i] + Y[i + 1]) / 2.
         return area
 
+    def Append(self, other):
+        """
+        Append points from another graph
+        """
+        orig_len = len(self)
+        self.Set(orig_len + len(other))
+        ipoint = orig_len
+        if hasattr(self, 'SetPointError'):
+            for point in other:
+                self.SetPoint(ipoint, point.x.value, point.y.value)
+                self.SetPointError(
+                    ipoint,
+                    point.x.error_low, point.x.error_hi,
+                    point.y.error_low, point.y.error_hi)
+                ipoint += 1
+        else:
+            for point in other:
+                self.SetPoint(ipoint, point.x.value, point.y.value)
+                ipoint += 1
+
 
 class _Graph2DBase(_GraphBase):
 
diff --git a/rootpy/tree/chain.py b/rootpy/tree/chain.py
--- a/rootpy/tree/chain.py
+++ b/rootpy/tree/chain.py
@@ -6,9 +6,11 @@
 import time
 
 from .. import log; log = log[__name__]
+from .. import QROOT
 from ..io import root_open, DoesNotExist
 from ..utils.extras import humanize_bytes
 from ..context import preserve_current_directory
+from ..plotting.graph import _GraphBase
 from ..extern.six import string_types
 from .filtering import EventFilterList
 
@@ -114,17 +116,20 @@ def Draw(self, *args, **kwargs):
                 # Make our own copy of the drawn histogram
                 output = self._tree.Draw(*args, **kwargs)
                 if output is not None:
-                    # Make it memory resident
                     output = output.Clone()
-                    output.SetDirectory(0)
+                    # Make it memory resident (histograms)
+                    if hasattr(output, 'SetDirectory'):
+                        output.SetDirectory(0)
             else:
                 newoutput = self._tree.Draw(*args, **kwargs)
                 if newoutput is not None:
-                    output += newoutput
+                    if isinstance(output, _GraphBase):
+                        output.Append(newoutput)
+                    else:  # histogram
+                        output += newoutput
         return output
 
-    def draw(self, *args, **kwargs):
-        return self.Draw(*args, **kwargs)
+    draw = Draw
 
     def __getattr__(self, attr):
         try:
@@ -141,7 +146,8 @@ def __contains__(self, branch):
 
     def __iter__(self):
         passed_events = 0
-        while True:
+        self.reset()
+        while self._rollover():
             entries = 0
             total_entries = float(self._tree.GetEntries())
             t1 = time.time()
@@ -171,18 +177,17 @@ def __iter__(self):
                 self._file.GetBytesRead(),
                 self._file.GetReadCalls()))
             self._total_events += entries
-            if not self._rollover():
-                break
         self._filters.finalize()
 
     def _rollover(self):
-        BaseTreeChain.reset(self)
         filename = self._next_file()
         if filename is None:
             return False
         log.info("current file: {0}".format(filename))
         try:
             with preserve_current_directory():
+                if self._file is not None:
+                    self._file.Close()
                 self._file = root_open(filename)
         except IOError:
             self._file = None
@@ -246,6 +251,15 @@ def __init__(self, name, files, **kwargs):
         self._files = files
         self.curr_file_idx = 0
         super(TreeChain, self).__init__(name, **kwargs)
+        self._tchain = QROOT.TChain(name)
+        for filename in self._files:
+            self._tchain.Add(filename)
+
+    def GetEntries(self, *args, **kwargs):
+        return self._tchain.GetEntries(*args, **kwargs)
+
+    def GetEntriesFast(self, *args, **kwargs):
+        return self._tchain.GetEntriesFast(*args, **kwargs)
 
     def reset(self):
         """
@@ -271,7 +285,14 @@ def _next_file(self):
 
 
 class TreeQueue(BaseTreeChain):
+    """
+    A chain of files in a multiprocessing Queue.
 
+    Note that asking for the number of files in the queue with len(treequeue)
+    can be unreliable. Also, methods not overridden by TreeQueue will always be
+    called on the current tree, so GetEntries will give you the number of
+    entries in the current tree.
+    """
     SENTINEL = None
 
     def __init__(self, name, files, **kwargs):


Test patch:
diff --git a/rootpy/tree/tests/test_tree.py b/rootpy/tree/tests/test_tree.py
--- a/rootpy/tree/tests/test_tree.py
+++ b/rootpy/tree/tests/test_tree.py
@@ -7,6 +7,7 @@
 from rootpy.io import root_open, TemporaryFile
 from rootpy.tree.treetypes import FloatCol, IntCol
 from rootpy.plotting import Hist, Hist2D, Hist3D
+from rootpy.plotting.graph import _GraphBase
 from rootpy import testdata
 from rootpy import stl
 
@@ -20,7 +21,7 @@
     from io import StringIO
 
 from nose.plugins.skip import SkipTest
-from nose.tools import (assert_raises, assert_almost_equal,
+from nose.tools import (assert_raises, assert_true, assert_almost_equal,
                         assert_equal, raises, with_setup)
 
 
@@ -51,11 +52,11 @@ class Event(ObjectA.prefix('a_') + ObjectB.prefix('b_')):
     return Event
 
 
-def create_tree():
+def create_tree(entries=100):
     f = TemporaryFile()
     tree = Tree("tree", model=create_model())
     # fill the tree
-    for i in range(1000):
+    for i in range(entries):
         assert_equal(tree.a_vect, LorentzVector(0, 0, 0, 0))
         random_vect = LorentzVector(
             gauss(.5, 1.),
@@ -90,9 +91,9 @@ def create_tree():
     FILE_PATHS.append(f.GetName())
 
 
-def create_chain():
+def create_chain(entries_per_tree=100):
     for i in range(3):
-        create_tree()
+        create_tree(entries=entries_per_tree)
 
 
 def cleanup():
@@ -135,9 +136,9 @@ def test_draw():
         tree.draw('a_x:a_y:a_z:b_x')
         tree.draw('a_x:a_y:a_z:b_x:b_y', options='para')
 
-        h1 = Hist(10, -1, 2, name='h1')
-        h2 = Hist2D(10, -1, 2, 10, -1, 2)
-        h3 = Hist3D(10, -1, 2, 10, -1, 2, 10, -1, 2)
+        h1 = Hist(10, -10, 10, name='h1')
+        h2 = Hist2D(10, -10, 10, 10, -10, 10)
+        h3 = Hist3D(10, -10, 10, 10, -10, 10, 10, -10, 10)
 
         # dimensionality does not match
         assert_raises(TypeError, tree.draw, 'a_x:a_y', hist=h1)
@@ -192,6 +193,24 @@ def test_draw():
         assert_equal(hist3.Integral(), hist1.Integral())
 
 
+@with_setup(create_chain, cleanup)
+def test_chain_iter():
+    if sys.version_info[0] >= 3:
+        raise SkipTest("Python 3 support not implemented")
+    chain = TreeChain('tree', FILE_PATHS)
+    assert_equal(len(chain), 3)  # 3 files
+    entries = 0
+    for entry in chain:
+        entries += 1
+    assert_equal(entries, 300)
+    entries = 0
+    for entry in chain:
+        entries += 1
+    assert_equal(entries, 300)
+    assert_equal(chain.GetEntries(), 300)
+    assert_equal(chain.GetEntriesFast(), 300)
+
+
 @with_setup(create_chain, cleanup)
 def test_chain_draw():
     if sys.version_info[0] >= 3:
@@ -200,12 +219,19 @@ def test_chain_draw():
     hist = Hist(100, 0, 1)
     chain.draw('a_x', hist=hist)
     assert_equal(hist.Integral() > 0, True)
+    assert_equal(hist.GetEntries(), 300)
 
     # check that Draw can be repeated
     hist2 = Hist(100, 0, 1)
     chain.draw('a_x', hist=hist2)
     assert_equal(hist.Integral(), hist2.Integral())
 
+    # draw into a graph
+    graph = chain.draw("a_x:a_y")
+    assert_true(isinstance(graph, _GraphBase))
+    assert_equal(len(graph), chain.GetEntries())
+    assert_equal(len(graph), 300)
+
 
 @with_setup(create_chain, cleanup)
 def test_chain_draw_hist_init_first():