Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_fast_grouped_custom_user_func_fail():
@symbolic_dispatch
def f(x):
return GroupByAgg.from_result(x.mean(), x)
gdf = data_default.groupby('g')
with pytest.raises(FunctionLookupError):
g_out = fast_mutate(gdf, result1 = f(_.x), result2 = _.x.mean() + 10)
from siuba.siu import Call, CallTreeLocal, str_to_getitem_call, Lazy, FunctionLookupError
# TODO: currently needed for select, but can we remove pandas?
from pandas import Series
import pandas as pd
from sqlalchemy.sql import schema
# TODO:
# - distinct
# - annotate functions using sel.prefix_with("\n/**/\n") ?
# Helpers ---------------------------------------------------------------------
class SqlFunctionLookupError(FunctionLookupError): pass
class CallListener:
"""Generic listener. Each exit is called on a node's copy."""
def enter(self, node):
args, kwargs = node.map_subcalls(self.enter)
return self.exit(node.__class__(node.func, *args, **kwargs))
def exit(self, node):
return node
class WindowReplacer(CallListener):
"""Call tree listener.
def create_local_call(self, name, prev_obj, cls, func_args = None, func_kwargs = None):
# need call attr name (arg[0].args[1])
# need call arg and kwargs
func_args = tuple() if func_args is None else func_args
func_kwargs = {} if func_kwargs is None else func_kwargs
try:
local_func = self.local[name]
except KeyError as err:
raise FunctionLookupError("Missing translation for function call: %s"% name)
if isclass(local_func) and issubclass(local_func, Exception):
raise local_func
return cls(
"__call__",
local_func,
prev_obj,
*func_args,
**func_kwargs
)
func = node(None)
# TODO: not robust at all, need class for singledispatch? unique attr flag?
if (hasattr(func, 'registry')
and hasattr(func, 'dispatch')
and self.dispatch_cls is not None
):
# allow custom functions that dispatch on dispatch_cls
f_for_cls = func.dispatch(self.dispatch_cls)
if (self.result_cls is None
or is_dispatch_func_subtype(f_for_cls, self.dispatch_cls, self.result_cls)
):
# matches return annotation type (or not required)
return node.__class__(f_for_cls)
raise FunctionLookupError(
"External function {name} can dispatch on the class {dispatch_cls}, but "
"must also have result annotation of (sub)type {result_cls}"
.format(
name = func.__name__,
dispatch_cls = self.dispatch_cls,
result_cls = self.result_cls
)
)
# doesn't raise an error so we can look in locals for now
# TODO: remove behavior, once all SQL dispatch funcs moved from locals
return self.generic_enter(node)
# determine dispatch class
cls_name = 'window' if window else 'aggregate'
dispatch_cls = self.dispatch_cls[cls_name]
call_shaper = CallTreeLocal(
funcs,
call_sub_attr = self.call_sub_attr,
dispatch_cls = dispatch_cls,
result_cls = self.result_cls
)
# raise informative error message if missing translation
try:
return call_shaper.enter(call)
except FunctionLookupError as err:
raise SqlFunctionLookupError.from_verb(
verb_name or "Unknown",
arg_name or "Unknown",
err,
short = True
)