def test_select_colmns(self):
ttree = Mock()
ttree.arrays = Mock(return_value=[1, 2, 3])
with patch("irishep.datasets.awkward_dataset.AwkwardDataset.__init__",
return_value=None) as awk_dataset:
d = UprootDataset("foo", ttree)
d2 = d.select_columns(["a", "b", "c"])
awk_dataset.assert_called_with("foo", [1, 2, 3])
self.assertIsInstance(d2, AwkwardDataset)
def read_files(self, dataset_name, files):
if len(files) > 1:
print("WARN: Uproot implementation doesn't work with multiple " +
"files in a dataset. Just reading the first file")
root = uproot.open(files[0])
dataset = UprootDataset(dataset_name, root["Events"])
return dataset
def test_execute_udf(self):
d = UprootDataset("foo", [1, 2, 3])
with self.assertRaises(NotImplementedError):
d.execute_udf(43)
def test_repartition(self):
d = UprootDataset("foo", [1, 2, 3])
with self.assertRaises(NotImplementedError):
d.repartition(43)
def test_show(self):
d = UprootDataset("foo", [1, 2, 3])
with self.assertRaises(NotImplementedError):
d.show()
def test_columns_with_types(self):
d = UprootDataset("foo", [1, 2, 3])
with self.assertRaises(NotImplementedError):
# noinspection PyStatementEffect
d.columns_with_types
def test_columns(self):
ttree = {b'a': 1, b'b': 2, b'c': 3}
d = UprootDataset("foo", ttree)
cols = d.columns
self.assertEqual(["a", "b", "c"], cols)
def test_count(self):
ttree = [1, 2, 3]
d = UprootDataset("foo", ttree)
self.assertEqual(3, d.count())
def test_init(self):
ttree = Mock()
d = UprootDataset("foo", ttree)
self.assertEqual(d.name, "foo")
self.assertEqual(d.ttree, ttree)