Repository: scikit-hep/awkward

Issue:
`with_cache` raises `NameError` if cache is given and layout already has one
Traceback
```pytb
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-55-18e9a13fc6a7> in <module>
----> 1 dataset = _from_hdf5_dataset("/tmp/dataset")

/mnt/app/texat/src/texat/awkward/io.ipynb in _from_hdf5_dataset(directory, lazy, h5py_options)
     82     partitions = sum([ak.partition.every(l) for l in layouts], start=[])
     83     layout = ak.partitioned(partitions, highlevel=False)
---> 84     cached = ak.with_cache(layout, cache={})
     85     return ak.Array(cached)
     86 

/opt/texat-venv/lib/python3.9/site-packages/awkward/operations/structure.py in with_cache(array, cache, highlevel, behavior)
   4193             return None
   4194 
-> 4195     out = ak._util.recursively_apply(
   4196         ak.operations.convert.to_layout(array), getfunction, pass_depth=False
   4197     )

/opt/texat-venv/lib/python3.9/site-packages/awkward/_util.py in recursively_apply(layout, getfunction, pass_depth, pass_user, pass_apply, user, keep_parameters, numpy_to_regular)
   1322             )
   1323 
-> 1324     return apply(layout, 1, user)
   1325 
   1326 

/opt/texat-venv/lib/python3.9/site-packages/awkward/_util.py in apply(layout, depth, user)
   1127         if isinstance(layout, ak.partition.PartitionedArray):
   1128             return ak.partition.IrregularlyPartitionedArray(
-> 1129                 [apply(x, depth, user) for x in layout.partitions]
   1130             )
   1131 

/opt/texat-venv/lib/python3.9/site-packages/awkward/_util.py in <listcomp>(.0)
   1127         if isinstance(layout, ak.partition.PartitionedArray):
   1128             return ak.partition.IrregularlyPartitionedArray(
-> 1129                 [apply(x, depth, user) for x in layout.partitions]
   1130             )
   1131 

/opt/texat-venv/lib/python3.9/site-packages/awkward/_util.py in apply(layout, depth, user)
   1120         custom = getfunction(layout, *args)
   1121         if callable(custom):
-> 1122             return custom()
   1123         else:
   1124             user = custom

/opt/texat-venv/lib/python3.9/site-packages/awkward/operations/structure.py in <lambda>()
   4185             return lambda: ak.layout.VirtualArray(
   4186                 layout.generator,
-> 4187                 newcache,
   4188                 layout.cache_key,
   4189                 layout.identities,

NameError: free variable 'newcache' referenced before assignment in enclosing scope
```


Discussion:

Resolving patch:
diff --git a/src/awkward/operations/structure.py b/src/awkward/operations/structure.py
--- a/src/awkward/operations/structure.py
+++ b/src/awkward/operations/structure.py
@@ -4168,6 +4168,16 @@ def with_cache(array, cache, highlevel=True, behavior=None):
 
     See #ak.virtual.
     """
+    if not highlevel:
+        raise NotImplementedError(
+            "ak.with_cache cannot allow highlevel=False because the only strong references\n"
+            "to caches are held by ak.Array objects; VirtualArrays only hold weak references,\n"
+            "which would go out of scope with this function. This will be fixed in Awkward 2.0,\n"
+            "when VirtualArrays are reimplemented in Python and can safely hold strong\n"
+            "references to caches.\n\n"
+            "For now, use highlevel=True and extract the layout from the output array."
+        )
+
     if cache == "new":
         hold_cache = ak._util.MappingProxy({})
         cache = ak.layout.ArrayCache(hold_cache)
@@ -4179,7 +4189,7 @@ def getfunction(layout):
         if isinstance(layout, ak.layout.VirtualArray):
             if cache is None:
                 newcache = layout.cache
-            elif layout.cache is None:
+            else:
                 newcache = cache
             return lambda: ak.layout.VirtualArray(
                 layout.generator,


Test patch:
diff --git a/tests/test_940-with-cache-already-exists.py b/tests/test_940-with-cache-already-exists.py
new file mode 100644
--- /dev/null
+++ b/tests/test_940-with-cache-already-exists.py
@@ -0,0 +1,28 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+from __future__ import absolute_import
+
+import pytest  # noqa: F401
+import numpy as np  # noqa: F401
+import awkward as ak  # noqa: F401
+
+
+def func():
+    return ak.Array([1, 2, 3, 4, 5])
+
+
+generator = ak.layout.ArrayGenerator(func, (), length=5)
+
+
+def test():
+    hold_cache = ak._util.MappingProxy({})
+    cache = ak.layout.ArrayCache(hold_cache)
+    layout = ak.layout.VirtualArray(generator, cache=cache)
+    array = ak.Array(layout)
+
+    cache_2 = {}
+    array_2 = ak.with_cache(array, cache_2)
+    ak.materialized(array_2)
+
+    assert len(cache_2) > 0
+    assert len(cache) == 0