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)