def test_find_root_no_root(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
a = ast_DataFrame(df)
r = _find_root_expr(ast.Num(n=10, ctx=ast.Load()), a)
assert r is None
def _term_to_ast(term: Union[int, str, DataFrame, Column, Callable],
parent_df: Optional[Union[DataFrame, Column]]) \
-> ast.AST:
'''Return an AST that represents the current term
Args:
term The term (int, string, float, DataFrame, Column, etc.)
Returns
'''
other_ast = None
if isinstance(term, int) or isinstance(term, float):
other_ast = ast.Num(n=term)
elif isinstance(term, str):
other_ast = ast.Str(s=term)
elif isinstance(term, DataFrame):
other_ast = ast_DataFrame(term)
elif isinstance(term, Column):
other_ast = ast_Column(term)
elif callable(term):
assert parent_df is not None, \
'Internal Error: Parent DF is required when creating a Callable'
assert isinstance(parent_df, DataFrame)
other_ast = ast_Callable(term, parent_df)
elif isinstance(term, tuple):
other_ast = ast.Tuple(
[_term_to_ast(item, parent_df) for item in cast(tuple, term)])
elif isinstance(term, list):
other_ast = ast.List(
[_term_to_ast(item, parent_df) for item in cast(list, term)])
else:
raise DataFrameTypeError("Do not know how to render a term "
f"of type '{type(term).__name__}'.")
return other_ast
def test_find_root_ast_df_simple(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
a = ast_DataFrame(df)
r = _find_root_expr(a, ast.Num(n=10, ctx=ast.Load()))
assert r is not None
assert r is a
def _lookup_dataframe(self, df: DataFrame) -> ast_DataFrame:
'''
See if a raw dataframe has already been tagged. If so, make sure
we return the same object.
'''
h = hash(str(df))
if h not in self._seen_datasources:
self._seen_datasources[h] = ast_DataFrame(df)
return self._seen_datasources[h]
def test_find_root_in_function(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
a = ast_DataFrame(df)
attr = ast.Attribute(value=a, attr='jets', ctx=ast.Load())
call = ast.Call(func=ast.Name(id='sin'), args=[attr], keywords=None)
r = _find_root_expr(call, attr)
assert r is attr
def test_find_root_ast_df_nested(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
a = ast_DataFrame(df)
attr = ast.Attribute(value=a, attr='jets', ctx=ast.Load())
r = _find_root_expr(attr, ast.Num(n=10, ctx=ast.Load()))
assert r is not None
assert r is a
def test_render_context_clone_resolved():
r1 = render_context()
r1._resolved[1] = ast_DataFrame(DataFrame())
r2 = render_context(r1)
assert 1 in r2._resolved
assert r2._resolved[1] is r1._resolved[1]
del r1._resolved[1]
assert 1 in r2._resolved
def test_render_context_clone_df():
r1 = render_context()
r1._seen_datasources[1] = ast_DataFrame(DataFrame())
r2 = render_context(r1)
assert 1 in r2._seen_datasources
assert r2._seen_datasources[1] is r1._seen_datasources[1]
del r1._seen_datasources[1]
assert 1 in r2._seen_datasources
def test_find_root_arg_not_right(servicex_ds):
f = ServiceXDatasetSource(servicex_ds)
df = xaod_table(f)
a = ast_DataFrame(df)
# df.ele.deltar(df.jets), with df.jets as the arg.
jets_attr = ast.Attribute(value=a, attr='jets', ctx=ast.Load())
eles_attr = ast.Attribute(value=a, attr='eles', ctx=ast.Load())
call = ast.Call(func=eles_attr, args=[jets_attr], keywords=None)
r = _find_root_expr(call, jets_attr)
assert r is a
def __init__(self, df: DataFrame):
DataFrame.__init__(self, expr=ast_DataFrame(df))
self.__p = df
def __init__(self,
df: Optional[DataFrame],
compound: Optional[op_base] = None) -> None:
DataFrame.__init__(self, expr=ast_DataFrame(df))
self._ref: op_base = compound if compound is not None else op_vec(self)
def test_df_none():
'Need to make sure the blank ctor is legal for deep copy reasons'
_ = ast_DataFrame()
def __init__(self,
df: DataFrame,
compound: Optional[op_base] = None) -> None:
DataFrame.__init__(self, ast_DataFrame(df))
self._base_df = df
self._ref: op_base = compound if compound is not None else op_vec(df)
def _as_ast(df: Union[DataFrame, Column]) -> Union[ast_DataFrame, ast_Column]:
if isinstance(df, DataFrame):
return ast_DataFrame(df)
else:
return ast_Column(df)
