Set Tree Operation

servicex/func_adl/func_adl_dataset.py

@@ -41,7 +41,6 @@
     Iterable,
     Dict,
 )
-from pydantic import typing
 from qastle import python_ast_to_text_ast
 
 from func_adl import EventDataset
@@ -99,29 +98,25 @@ def __init__(
     def set_provided_qastle(self, qastle: str):
         self.provided_qastle = qastle
 
-    def clone_with_new_ast(self, new_ast: ast.AST, new_type: typing.Any):
+    def __deepcopy__(self, memo):
         """
-        Override the method from ObjectStream - We need to be careful because the query
-        cache is a tinyDB database that holds an open file pointer. We are not allowed
-        to clone an open file handle, so for this property we will copy by reference
-        and share it between the clones. Turns out ast class is also picky about copies,
-        so we set that explicitly.
-
-        :param new_ast:
-        :param new_type:
-        :return:
+        Customize deepcopy behavior for this class.
+        We need to be careful because the query cache is a tinyDB database that holds an
+        open file pointer. We are not allowed to clone an open file handle, so for this
+        property we will copy by reference and share it between the clones
         """
-        clone = copy.copy(self)
+        cls = self.__class__
+        obj = cls.__new__(cls)
+
+        memo[id(self)] = obj
+
         for attr, value in vars(self).items():
             if type(value) == QueryCache:
-                setattr(clone, attr, value)
-            elif attr == "_q_ast":
-                setattr(clone, attr, new_ast)
+                setattr(obj, attr, value)
             else:
-                setattr(clone, attr, copy.deepcopy(value))
+                setattr(obj, attr, copy.deepcopy(value, memo))
 
-        clone._item_type = new_type
-        return clone
+        return obj
 
     def SelectMany(
         self, func: Union[str, ast.Lambda, Callable[[T], Iterable[S]]]
@@ -216,6 +211,20 @@ def QMetaData(self, metadata: Dict[str, Any]) -> FuncADLDataset[T]:
         """
         return super().QMetaData(metadata)
 
+    def set_tree(self, tree_name: str) -> FuncADLDataset[T]:
+        r"""Set the tree name for the query.
+        Args:
+            tree_name (str): Name of the tree to use for the query
+
+        Returns:
+            The Dataset with the tree appended to the first call object
+        """
+
+        clone = self.clone_with_new_ast(copy.deepcopy(self.query_ast), self._item_type)
+        clone._q_ast.args[0].args.append(ast.Str(s="bogus.root"))
+        clone._q_ast.args[0].args.append(ast.Str(s=tree_name))
+        return clone
+
     def generate_qastle(self, a: ast.AST) -> str:
         r"""Generate the qastle from the ast of the query.