Merge branch 'master' into issue-431

README.rst

@@ -77,6 +77,22 @@ Optional dependencies:
 
 **Reminder: you do not need C++ ROOT to run uproot.**
 
+Testing and development requirements
+------------------------------------
+
+The packages need to run the test suite can be installed via ``pip install
+"uproot[testing]"``. Here is an example how to set up a development
+
+.. code-block:: bash
+
+    git clone https://github.com/scikit-hep/uproot.git  # alternatively your own fork
+    cd uproot
+    python -m venv venv  # create a virtual environment in the folder venv
+    . venv/bin/activate  # activate the Python environment
+    pip install -e ".[testing]"   # installs uproot in editable mode with all packages required for testing
+    pytest               # run the test suite
+
+
 .. inclusion-marker-3-do-not-remove
 
 Questions

setup.py

@@ -144,7 +144,9 @@ def get_description():
       python_requires = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*",
       install_requires = ["numpy>=1.13.1", "awkward>=0.12.0,<1.0", "uproot-methods>=0.7.0", "cachetools"],
       setup_requires = ["pytest-runner"],
-      tests_require = ["pytest>=3.9", "pkgconfig", "lz4", "zstandard", 'backports.lzma;python_version<"3.3"', "xxhash", "mock", "requests"],
+      extras_require = {
+          'testing': ["pytest>=3.9", "pkgconfig", "lz4", "zstandard", 'backports.lzma;python_version<"3.3"', "xxhash", "mock", "requests"],
+      },
       classifiers = [
           "Development Status :: 5 - Production/Stable",
           "Intended Audience :: Developers",

tests/test_issues.py

@@ -396,3 +396,24 @@ def test_issue431(self):
         file = uproot.open("tests/samples/issue431.root")
         head = file["Head"]
         assert head._map_3c_string_2c_string_3e_ == {b'DAQ': b'394', b'PDF': b'4      58', b'XSecFile': b'', b'can': b'0 1027 888.4', b'can_user': b'0.00 1027.00  888.40', b'coord_origin': b'0 0 0', b'cut_in': b'0 0 0 0', b'cut_nu': b'100 1e+08 -1 1', b'cut_primary': b'0 0 0 0', b'cut_seamuon': b'0 0 0 0', b'decay': b'doesnt happen', b'detector': b'NOT', b'drawing': b'Volume', b'end_event': b'', b'genhencut': b'2000 0', b'genvol': b'0 1027 888.4 2.649e+09 100000', b'kcut': b'2', b'livetime': b'0 0', b'model': b'1       2       0       1      12', b'muon_desc_file': b'', b'ngen': b'0.1000E+06', b'norma': b'0 0', b'nuflux': b'0       3       0 0.500E+00 0.000E+00 0.100E+01 0.300E+01', b'physics': b'GENHEN 7.2-220514 181116 1138', b'seed': b'GENHEN 3  305765867         0         0', b'simul': b'JSirene 11012 11/17/18 07', b'sourcemode': b'diffuse', b'spectrum': b'-1.4', b'start_run': b'1', b'target': b'isoscalar', b'usedetfile': b'false', b'xlat_user': b'0.63297', b'xparam': b'OFF', b'zed_user': b'0.00 3450.00'}
+
+    def test_issue434(self):
+        f = uproot.open("tests/samples/issue434.root")
+        fromdtype = [("pmt", "u1"), ("tdc", "<u4"), ("tot", "u1")]
+        todtype = [("pmt", "u1"), ("tdc", ">u4"), ("tot", "u1")]
+        tree = f[b'KM3NET_TIMESLICE_L1'][b'KM3NETDAQ::JDAQTimeslice']
+        superframes = tree[b'vector<KM3NETDAQ::JDAQSuperFrame>']
+        hits_buffer = superframes[b'vector<KM3NETDAQ::JDAQSuperFrame>.buffer']
+        hits = hits_buffer.lazyarray(
+                uproot.asjagged(
+                    uproot.astable(
+                        uproot.asdtype(fromdtype, todtype)), skipbytes=6))
+        assert 486480 == hits['tdc'][0][0]
+
+    def test_issue438_accessing_memory_mapped_objects_outside_of_context_raises(self):
+        with uproot.open("tests/samples/issue434.root") as f:
+            a = f['KM3NET_EVENT']['KM3NET_EVENT']['KM3NETDAQ::JDAQPreamble'].array()
+            b = f['KM3NET_EVENT']['KM3NET_EVENT']['KM3NETDAQ::JDAQPreamble'].lazyarray()
+        assert 4 == len(a[0])
+        with pytest.raises(IOError):
+            len(b[0])

tests/test_write.py

@@ -1815,3 +1815,43 @@ def test_large_compress_uproot(tmp_path):
     f = uproot.open(filename)
     assert f["a"] == ("a"*((2**24) + 2000)).encode("utf-8")
     assert f["b"] == ("b"*((2**24) + 10)).encode("utf-8")
+
+def test_tree_twodim(tmp_path):
+    filename = join(str(tmp_path), "example.root")
+
+    a = numpy.array([[0, 1, 2, 3],
+                     [3, 4, 5, 6]])
+
+    with uproot.recreate(filename, compression=None) as f:
+        f["t"] = uproot.newtree({"branch": uproot.newbranch(numpy.dtype(">i4"), shape=a.shape)})
+        f["t"].extend({"branch": a})
+
+    f = ROOT.TFile.Open(filename)
+    tree = f.Get("t")
+    rdf = ROOT.RDataFrame(tree).AsNumpy()["branch"]
+    for i in range(0, 2):
+        for j in range(0, 4):
+            assert a[i][j] == rdf[i][j]
+
+def test_tree_threedim(tmp_path):
+    filename = join(str(tmp_path), "example.root")
+
+    a = numpy.array([[[0, 1, 2, 3],
+                      [3, 4, 5, 6],
+                      [90, 91, 91, 92]],
+                     [[10, 11, 12, 13],
+                      [13, 14, 15, 16],
+                      [190, 191, 191, 192]]])
+
+    with uproot.recreate(filename, compression=None) as f:
+        f["t"] = uproot.newtree({"branch": uproot.newbranch(numpy.dtype(">i4"), shape=a.shape)})
+        f["t"].extend({"branch": a})
+
+    f = ROOT.TFile.Open(filename)
+    tree = f.Get("t")
+    rdf = ROOT.RDataFrame(tree).AsNumpy()["branch"]
+    for i in range(2):
+        test = numpy.array(rdf[i]).reshape(3, 4)
+        for j in range(3):
+            for k in range(4):
+                assert a[i][j][k] == test[j][k]

uproot/interp/numerical.py

@@ -17,6 +17,9 @@
 else:
     string_types = (str, bytes)
 
+BYTEORDER_INDICATORS = (">", "<", "=", "|", b">", b"<", b"=", b"|")
+
+
 def _dtypeshape(obj):
     out = ()
     while obj.subdtype is not None:
@@ -85,7 +88,9 @@ class asdtype(_asnumeric):
     def __init__(self, fromdtype, todtype=None):
         if isinstance(fromdtype, self.awkward.numpy.dtype):
             self.fromdtype = fromdtype
-        elif isinstance(fromdtype, string_types) and len(fromdtype) > 0 and fromdtype[0] in (">", "<", "=", "|", b">", b"<", b"=", b"|"):
+        elif isinstance(fromdtype, string_types) and len(fromdtype) > 0 and fromdtype[0] in BYTEORDER_INDICATORS:
+            self.fromdtype = self.awkward.numpy.dtype(fromdtype)
+        elif isinstance(fromdtype, list) and any(e[1][0] in BYTEORDER_INDICATORS for e in fromdtype):
             self.fromdtype = self.awkward.numpy.dtype(fromdtype)
         else:
             self.fromdtype = self.awkward.numpy.dtype(fromdtype).newbyteorder(">")
@@ -94,7 +99,9 @@ def __init__(self, fromdtype, todtype=None):
             self.todtype = self.fromdtype.newbyteorder("=")
         elif isinstance(todtype, self.awkward.numpy.dtype):
             self.todtype = todtype
-        elif isinstance(todtype, string_types) and len(todtype) > 0 and todtype[0] in (">", "<", "=", "|", b">", b"<", b"=", b"|"):
+        elif isinstance(todtype, string_types) and len(todtype) > 0 and todtype[0] in BYTEORDER_INDICATORS:
+            self.todtype = self.awkward.numpy.dtype(todtype)
+        elif isinstance(todtype, list) and any(e[1][0] in BYTEORDER_INDICATORS for e in todtype):
             self.todtype = self.awkward.numpy.dtype(todtype)
         else:
             self.todtype = self.awkward.numpy.dtype(todtype).newbyteorder("=")

uproot/rootio.py

@@ -362,6 +362,9 @@ def get(self, name, cycle=None):
             else:
                 raise _KeyError("not found: {0} with cycle {1}\n in file: {2}".format(repr(name), cycle, self._context.sourcepath))
 
+    def close(self):
+        self._context.source.close()
+
     def __contains__(self, name):
         try:
             self.get(name)
@@ -374,7 +377,7 @@ def __enter__(self, *args, **kwds):
         return self
 
     def __exit__(self, *args, **kwds):
-        pass
+        self.close()
 
 class _KeyError(KeyError):
     def __str__(self):

uproot/source/chunked.py

@@ -46,6 +46,10 @@ def __del__(self):
     def _read(self, chunkindex):
         raise NotImplementedError
 
+    def close(self):
+        super(ChunkedSource, self).close()
+        self.cache.clear()
+
     def dismiss(self):
         if self._futures is not None:
             for future in self._futures.values():

uproot/source/memmap.py

@@ -19,12 +19,19 @@ class MemmapSource(uproot.source.source.Source):
     def __init__(self, path):
         self.path = os.path.expanduser(path)
         self._source = numpy.memmap(self.path, dtype=numpy.uint8, mode="r")
+        self.closed = False
+
+    @property
+    def source(self):
+        if self.closed:
+            raise IOError("The file handler has already been closed.")
+        return self._source
 
     def parent(self):
         return self
 
     def size(self):
-        return len(self._source)
+        return len(self.source)
 
     def threadlocal(self):
         return self
@@ -33,17 +40,18 @@ def dismiss(self):
         pass
 
     def close(self):
-        self._source._mmap.close()
+        self.source._mmap.close()
+        self.closed = True
 
     def data(self, start, stop, dtype=None):
         # assert start >= 0
         # assert stop >= 0
         # assert stop >= start
 
-        if stop > len(self._source):
-            raise IndexError("indexes {0}:{1} are beyond the end of data source {2}".format(len(self._source), stop, repr(self.path)))
+        if stop > len(self.source):
+            raise IndexError("indexes {0}:{1} are beyond the end of data source {2}".format(len(self.source), stop, repr(self.path)))
 
         if dtype is None:
-            return self._source[start:stop]
+            return self.source[start:stop]
         else:
-            return self._source[start:stop].view(dtype)
+            return self.source[start:stop].view(dtype)

uproot/write/TKey.py

@@ -10,10 +10,10 @@
 import uproot.write.util
 
 class BasketKey(object):
-    def __init__(self, fName, fNevBuf, fNevBufSize, fObjlen=0, fSeekKey=100, fSeekPdir=0, fBufferSize=0):
+    def __init__(self, fName, fTitle, fNevBuf, fNevBufSize, fObjlen=0, fSeekKey=100, fSeekPdir=0, fBufferSize=0):
         self.fClassName = b"TBasket"
         self.fName = fName
-        self.fTitle = b"t"
+        self.fTitle = fTitle
 
         self.fObjlen = fObjlen
         self.fSeekKey = fSeekKey

uproot/write/objects/TTree.py

@@ -21,10 +21,11 @@
 
 class newbranch(object):
 
-    def __init__(self, type, title="", **options):
+    def __init__(self, type, title="", shape=(1,), **options):
         self.name = ""
         self.type = type
         self.title = title
+        self.shape = shape
         if "compression" in options:
             self.compression = options["compression"]
             del options["compression"]
@@ -86,12 +87,12 @@ def extend(self, branchdict):
         for key, value in branchdict.items():
             branchdict[key] = numpy.array(value, dtype=self._branches[key]._branch.type, copy=False)
 
-            #Temporary place holder
-            if branchdict[key].ndim != 1:
-                raise NotImplementedError("Multi dimensional array support is coming soon")
-
         for key, value in branchdict.items():
-            self._branches[key].newbasket(value)
+            if value.ndim == 1:
+                self._branches[key].newbasket(value)
+            else:
+                for i in range(0, value.shape[0]):
+                    self._branches[key].newbasket(value[i], i + 1)
 
     @property
     def name(self):
@@ -216,7 +217,7 @@ def revertstring(string):
             return string
 
     _tree_size = struct.Struct(">qqq")
-    def newbasket(self, items):
+    def newbasket(self, items, multidim=None):
         if len(items) == 0:
             return
 
@@ -280,20 +281,22 @@ def newbasket(self, items):
             self._treelvl1._tree = tree
             self._treelvl1._branches = temp_branches
 
-        self._branch.fields["_fEntries"] += len(items)
+        if multidim is None:
+            self._branch.fields["_fEntries"] += len(items)
+            self._branch.fields["_fEntryNumber"] += len(items)
+        else:
+            self._branch.fields["_fEntries"] = multidim
+            self._branch.fields["_fEntryNumber"] = multidim
         self._branch.fields["_fBasketEntry"][self._branch.fields["_fWriteBasket"]] = self._branch.fields["_fEntries"]
-        self._branch.fields["_fEntryNumber"] += len(items)
         basketdata = numpy.array(items, dtype=self._branch.type, copy=False)
 
-        if basketdata.ndim != 1:
-            raise NotImplementedError("Multi dimensional array support is coming soon")
-
         givenbytes = basketdata.tostring()
         cursor = uproot.write.sink.cursor.Cursor(self._branch.file._fSeekFree)
         self._branch.fields["_fBasketSeek"][self._branch.fields["_fWriteBasket"] - 1] = cursor.index
         key = BasketKey(fName=self._branch.name,
-                        fNevBuf=len(items),
-                        fNevBufSize=numpy.dtype(self._branch.type).itemsize,
+                        fTitle=self._treelvl1._tree.write_key.fName,
+                        fNevBuf=1 if multidim else len(items),
+                        fNevBufSize=numpy.dtype(self._branch.type).itemsize*len(items) if multidim else numpy.dtype(self._branch.type).itemsize,
                         fSeekKey=copy(self._branch.file._fSeekFree),
                         fSeekPdir=copy(self._branch.file._fBEGIN),
                         fBufferSize=32000)
@@ -681,6 +684,7 @@ class TBranchImpl(object):
     def __init__(self, name, branchobj, compression, file):
         self.name = _bytesid(name)
         self.type = numpy.dtype(branchobj.type).newbyteorder(">")
+        self.shape = branchobj.shape
         self.compression = compression
         self.util = None
         self.keycursor = None
@@ -708,6 +712,12 @@ def __init__(self, name, branchobj, compression, file):
                        "_fEntryNumber": 0,
                        "_fBaskets": b'@\x00\x00\x1d\x00\x03\x00\x01\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'}
 
+        # For multidimensional arrays
+        if len(self.shape) > 1:
+            array_pad = b""
+            for i in range(1, len(self.shape)):
+                array_pad += b"[" + str(self.shape[i]).encode("utf-8") + b"]"
+
         # TODO: Fix else condition to not always return NotImplementedError
         if self.type == "int8":
             title_pad = b"/B"
@@ -735,10 +745,18 @@ def __init__(self, name, branchobj, compression, file):
 
         if branchobj.title == "":
             self.title = _bytesid(name)
-            self.nametitle = self.title + title_pad
+            if len(self.shape) > 1:
+                self.nametitle = self.title + array_pad + title_pad
+                self.arraytitle = self.title + array_pad
+            else:
+                self.nametitle = self.title + title_pad
         else:
             self.title = _bytesid(branchobj.title)
-            self.nametitle = self.title
+            if len(self.shape) > 1:
+                self.nametitle = self.title + array_pad
+                self.arraytitle = self.title + array_pad
+            else:
+                self.nametitle = self.title
 
         self.fields["_fBasketBytes"] = numpy.array(self.fields["_fBasketBytes"], dtype=">i4", copy=False)
         self.fields["_fBasketEntry"] = numpy.array(self.fields["_fBasketEntry"], dtype=">i8", copy=False)
@@ -791,12 +809,15 @@ def put_tleaf(self, cursor):
         cursor.skip(self._format_cntvers.size)
         vers = 2
         fLen = 1
+        if len(self.shape) > 1:
+            for i in range(1, len(self.shape)):
+                fLen = fLen*self.shape[i]
         fLenType = numpy.dtype(self.type).itemsize
         fOffset = 0
         fIsRange = False
         fIsUnsigned = False
         fLeafCount = None
-        buff = (self.put_tnamed(cursor, self.name, self.title) +
+        buff = (self.put_tnamed(cursor, self.name, self.arraytitle if len(self.shape)>1 else self.title) +
                 cursor.put_fields(self._format_tleaf1, fLen, fLenType, fOffset, fIsRange, fIsUnsigned) +
                 self.util.put_objany(cursor, (fLeafCount, "TLeaf"), self.keycursor))
         length = len(buff) + self._format_cntvers.size