Repository: iris-hep/func_adl

Issue:
EventDataset can have as many arguments as it wants
Others might want to change the meaning (like uproot) of the arguments. Make sure that is possible in the EventDataset AST. The current `_extract_dataset_info` restricts the requirements to one.

Discussion:

Patch:
diff --git a/func_adl/event_dataset.py b/func_adl/event_dataset.py
index 8915d4c..34d7a69 100644
--- a/func_adl/event_dataset.py
+++ b/func_adl/event_dataset.py
@@ -83,7 +83,6 @@ def _extract_dataset_info(ds_call: ast.Call) -> EventDataset:
     args = cast(List[ast.AST], ds_call.args)
 
     # List should be strings
-    assert len(args) == 1
     return ast.literal_eval(args[0])
 
 


Test patch:
diff --git a/tests/test_event_dataset.py b/tests/test_event_dataset.py
index d413951..806ae84 100644
--- a/tests/test_event_dataset.py
+++ b/tests/test_event_dataset.py
@@ -17,6 +17,15 @@ class my_event(EventDataset):
         return 10
 
 
+class my_event_extra_args(EventDataset):
+    def __init__(self):
+        super().__init__()
+        self._ast.args.append(ast.Str(s='hi'))
+
+    async def execute_result_async(self, a: ast.AST) -> Any:
+        return 10
+
+
 def test_can_create():
     my_event()
 
@@ -32,6 +41,12 @@ def test_find_event_inside():
     assert find_ed_in_ast(add) is e
 
 
+def test_find_event_with_more_ast():
+    e = my_event_extra_args()
+    add = ast.BinOp(ast.Num(5), ast.Add, e._ast)
+    assert find_ed_in_ast(add) is e
+
+
 def test_string_rep():
     e = my_event()
     assert str(e) == "my_event"