Calls to AsXXX with column names as single items or single arrays should be the same (#99)

* If you call `AsAwkwardArray(['jetpt'])` or `AsAwkwardArray('jetpt')`, the produced `ast` will be identical.
* This helps with caching down the line.

Fixes #55

Author: gordonwatts

func_adl/object_stream.py

@@ -231,6 +231,9 @@ def AsPandasDF(
 
         """
         # To get Pandas use the ResultPandasDF function call.
+        if isinstance(columns, str):
+            columns = [columns]
+
         return ObjectStream[ReturnedDataPlaceHolder](
             function_call("ResultPandasDF", [self._q_ast, as_ast(columns)])
         )
@@ -262,6 +265,9 @@ def AsROOTTTree(
             dataset.  The order of the files back is consistent for different queries on the same
             dataset.
         """
+        if isinstance(columns, str):
+            columns = [columns]
+
         return ObjectStream[ReturnedDataPlaceHolder](
             function_call(
                 "ResultTTree", [self._q_ast, as_ast(columns), as_ast(treename), as_ast(filename)]
@@ -296,6 +302,9 @@ def AsParquetFiles(
             result. The order of the files back is consistent for different queries on the same
             dataset.
         """
+        if isinstance(columns, str):
+            columns = [columns]
+
         return ObjectStream[ReturnedDataPlaceHolder](
             function_call("ResultParquet", [self._q_ast, as_ast(columns), as_ast(filename)])
         )
@@ -316,6 +325,9 @@ def AsAwkwardArray(
 
             An `ObjectStream` with the `awkward` array data as its one and only element.
         """
+        if isinstance(columns, str):
+            columns = [columns]
+
         return ObjectStream[ReturnedDataPlaceHolder](
             function_call("ResultAwkwardArray", [self._q_ast, as_ast(columns)])
         )

tests/test_object_stream.py

@@ -75,6 +75,25 @@ def test_simple_query():
     assert isinstance(r, ast.AST)
 
 
+def test_two_simple_query():
+    r1 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsROOTTTree("junk.root", "analysis", "jetPT")
+        .value()
+    )
+    r2 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsROOTTTree("junk.root", "analysis", ["jetPT"])
+        .value()
+    )
+
+    assert ast.dump(r1) == ast.dump(r2)
+
+
 def test_with_types():
     r1 = my_event_with_type().SelectMany(lambda e: e.Jets("jets"))
     r = r1.Select(lambda j: j.eta()).value()
@@ -115,6 +134,25 @@ def test_simple_query_parquet():
     assert isinstance(r, ast.AST)
 
 
+def test_two_simple_query_parquet():
+    r1 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsParquetFiles("junk.root", "jetPT")
+        .value()
+    )
+    r2 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsParquetFiles("junk.root", ["jetPT"])
+        .value()
+    )
+
+    assert ast.dump(r1) == ast.dump(r2)
+
+
 def test_simple_query_panda():
     r = (
         my_event()
@@ -126,6 +164,24 @@ def test_simple_query_panda():
     assert isinstance(r, ast.AST)
 
 
+def test_two_imple_query_panda():
+    r1 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsPandasDF(["analysis"])
+        .value()
+    )
+    r2 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsPandasDF(["analysis"])
+        .value()
+    )
+    assert ast.dump(r1) == ast.dump(r2)
+
+
 def test_simple_query_awkward():
     r = (
         my_event()
@@ -137,6 +193,25 @@ def test_simple_query_awkward():
     assert isinstance(r, ast.AST)
 
 
+def test_two_similar_query_awkward():
+    r1 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsAwkwardArray(["analysis"])
+        .value()
+    )
+    r2 = (
+        my_event()
+        .SelectMany("lambda e: e.jets()")
+        .Select("lambda j: j.pT()")
+        .AsAwkwardArray("analysis")
+        .value()
+    )
+
+    assert ast.dump(r1) == ast.dump(r2)
+
+
 def test_metadata():
     r = (
         my_event()