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)