Merge pull request #67 from chrisburr/fix-60

Fix #60

Author: chrisburr

root_pandas/readwrite.py

@@ -199,11 +199,19 @@ def read_root(paths, key=None, columns=None, ignore=None, chunksize=None, where=
 
     if not isinstance(paths, list):
         paths = [paths]
-    # Use a single file to search for trees and branches
-    seed_path = paths[0]
+    # Use a single file to search for trees and branches, ensuring the key exists
+    for seed_path in paths:
+        trees = list_trees(seed_path)
+        if key and key not in trees:
+            continue
+        break
+    else:
+        if key:
+            raise OSError('{} not found in any of the given paths'.format(key))
+        else:
+            raise OSError('No trees found in any of the given paths')
 
     if not key:
-        trees = list_trees(seed_path)
         if len(trees) == 1:
             key = trees[0]
         elif len(trees) == 0:

tests/test.py

@@ -1,5 +1,5 @@
 import pandas as pd
-from root_pandas import read_root
+from root_pandas import read_root, to_root
 from root_numpy import list_branches
 from root_numpy import array2root
 from pandas.util.testing import assert_frame_equal
@@ -311,3 +311,13 @@ def test_brace_pattern_in_columns():
                        reference_df[['var02', 'var03', 'var11', 'var13']])
 
     os.remove('tmp.root')
+
+
+def test_detect_branches_first_missing():
+    df = pd.DataFrame({'a': list(range(10)), 'b': list(range(10))})
+    to_root(df, 'tmp_1.root', 'my_tree_1')
+    to_root(df, 'tmp_2.root', 'my_tree')
+    read_df = read_root(['tmp_1.root', 'tmp_2.root'], 'my_tree', warn_missing_tree=True)
+    assert_frame_equal(df, read_df)
+    os.remove('tmp_1.root')
+    os.remove('tmp_2.root')

tests/test_issues.py

@@ -4,6 +4,16 @@
 import root_pandas
 
 
+def test_issue_60():
+    df = pd.DataFrame({'a': list(range(10)), 'b': list(range(10))})
+    root_pandas.to_root(df, 'tmp_1.root', 'my_tree_1')
+    root_pandas.to_root(df, 'tmp_2.root', 'my_tree')
+    result = root_pandas.read_root(['tmp_1.root', 'tmp_2.root'], 'my_tree', warn_missing_tree=True)
+    assert len(result) == 10
+    os.remove('tmp_1.root')
+    os.remove('tmp_2.root')
+
+
 def test_issue_63():
     df = pd.DataFrame({'a': [], 'b': []})
     root_pandas.to_root(df, 'tmp_1.root', 'my_tree')