Reikriam

jagged i8 test

May 1st, 2020
124
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.25 KB | None | 0 0
  1. def test_jagged_i8(tmp_path):
  2. filename = join(str(tmp_path), "example.root")
  3.  
  4. a = awkward.fromiter([[0],
  5. [1, 2],
  6. [10, 11, 12]])
  7.  
  8. with uproot.recreate(filename, compression=None) as f:
  9. f["t"] = uproot.newtree({"branch": uproot.newbranch(numpy.dtype(">i8"), dependence="n")})
  10. f["t"].extend({"branch": a, "n": [1, 2, 3]})
  11.  
  12. ROOT.gInterpreter.Declare("""
  13. void readnewbasket(Long64_t* arr, char* filename) {
  14. TFile *f = new TFile("jaggedint.root");
  15. Long64_t x;
  16. Int_t num;
  17. Int_t last_num = 0;
  18. auto tree = f->Get<TTree>("t");
  19. auto n = tree->GetBranch("n");
  20. auto branch = tree->GetBranch("branch");
  21. n->SetAddress(&num);
  22. branch->SetAddress(&x);
  23. for (int i=0; i<tree->GetEntries(); i++) {
  24. tree->GetEvent(i);
  25. for (int j=0; j<num; j++) {
  26. *(arr + (i*last_num) + j) = x+j;
  27. }
  28. last_num = num;
  29. }
  30. }""")
  31.  
  32. testa = awkward.fromiter([[0], [0, 0], [0, 0, 0]])
  33. ROOT.readnewbasket(testa, filename)
  34. for i in range(3):
  35. assert (a[i] == testa[i])
Add Comment
Please, Sign In to add comment