def testWithFiles(self):
foo = temp.AutoTempFilePath(suffix="foo")
bar = temp.AutoTempFilePath(suffix="bar")
baz = temp.AutoTempFilePath(suffix="baz")
with context.MultiContext([foo, bar, baz]) as filepaths:
self.assertLen(filepaths, 3)
self.assertEndsWith(filepaths[0], "foo")
self.assertEndsWith(filepaths[1], "bar")
self.assertEndsWith(filepaths[2], "baz")
wbopen = functools.partial(io.open, mode="wb")
with context.MultiContext(map(wbopen, filepaths)) as filedescs:
self.assertLen(filedescs, 3)
filedescs[0].write(b"FOO")
filedescs[1].write(b"BAR")
filedescs[2].write(b"BAZ")
# At this point all three files should be correctly written, closed and
# ready for reading.
rbopen = functools.partial(io.open, mode="rb")
with context.MultiContext(map(rbopen, filepaths)) as filedescs:
self.assertLen(filedescs, 3)
self.assertEqual(filedescs[0].read(), b"FOO")
self.assertEqual(filedescs[1].read(), b"BAR")
self.assertEqual(filedescs[2].read(), b"BAZ")
def testWithNulls(self):
foo = context.NullContext("foo")
bar = context.NullContext("bar")
baz = context.NullContext("baz")
with context.MultiContext([foo, bar, baz]) as names:
self.assertEqual(names, ["foo", "bar", "baz"])
def VFSMultiOpen(pathspecs, progress_callback=None):
"""Opens multiple files specified by given path-specs.
See documentation for `VFSOpen` for more information.
Args:
pathspecs: A list of pathspec instances of files to open.
progress_callback: A callback function to call to notify about progress
Returns:
A context manager yielding file-like objects.
"""
precondition.AssertIterableType(pathspecs, rdf_paths.PathSpec)
vfs_open = functools.partial(VFSOpen, progress_callback=progress_callback)
return context.MultiContext(list(map(vfs_open, pathspecs)))
def testEmpty(self):
with context.MultiContext([]) as values:
self.assertEqual(values, [])
def MultiOpenAff4File(flow_obj, pathspecs):
aff4_paths = [_.AFF4Path(flow_obj.client_urn) for _ in pathspecs]
fileopens = aff4.FACTORY.MultiOpenOrdered(aff4_paths, token=flow_obj.token)
return context.MultiContext(fileopens)
