diff --git a/dovpanda/__init__.py b/dovpanda/__init__.py index 21325cd..d0f0ff6 100644 --- a/dovpanda/__init__.py +++ b/dovpanda/__init__.py @@ -1,10 +1,20 @@ import sys from dovpanda.core import ledger + if 'pandas' in sys.modules.keys(): ledger.register_hooks() else: ledger.tell('Pandas not imported') + def set_output(tell_method): - ledger.set_output(tell_method) \ No newline at end of file + ledger.set_output(tell_method) + + +def ignore_hook(func_name): + ledger.ignore_hook(func_name) + + +def reset_ignores(): + ledger.reset_ignores() diff --git a/dovpanda/base.py b/dovpanda/base.py index 8372b2f..830a235 100644 --- a/dovpanda/base.py +++ b/dovpanda/base.py @@ -3,6 +3,7 @@ import sys import traceback from collections import defaultdict +from itertools import chain import pandas @@ -13,6 +14,10 @@ pass +class HookFuncNotExists(Exception): + pass + + class Teller: def __init__(self): self.message = None @@ -66,10 +71,11 @@ class Ledger: def __init__(self): self.hooks = defaultdict(list) self.teller = Teller() + self.ignored_hooks = set() - def replace(self, original, hooks: tuple): + def replace(self, original, func_hooks: tuple): g = rgetattr(sys.modules['pandas'], original) - rsetattr(sys.modules['pandas'], original, attach_hooks(g, hooks)) + rsetattr(sys.modules['pandas'], original, self.attach_hooks(g, func_hooks)) def add_hook(self, original, hook_type='pre'): accepted_hooks = ['pre', 'post'] @@ -90,6 +96,38 @@ def tell(self, *args, **kwargs): def set_output(self, output): self.teller.set_output(output) + def ignore_hook(self, func_name): + hooks_names = [func[0].__name__ for func in chain.from_iterable([hooks for hooks in self.hooks.values()])] + if func_name in hooks_names: + self.ignored_hooks.update({func_name}) + else: + raise HookFuncNotExists( + f'Hook with function name: {func_name} not exists.\nNeed to choose one of: {hooks_names}') + + def reset_ignores(self): + self.ignored_hooks = set() + + def attach_hooks(self, f, func_hooks): + pres = [hook_function for (hook_function, hook_type) in func_hooks if hook_type.lower().startswith('pre')] + posts = [hook_function for (hook_function, hook_type) in func_hooks if hook_type.lower().startswith('post')] + + @functools.wraps(f) + def run(*args, **kwargs): + caller = traceback.extract_stack()[-2].filename + if caller.startswith(PANDAS_DIR): + ret = f(*args, **kwargs) + else: + for pre in pres: + if pre.__name__ not in self.ignored_hooks: + pre(*args, **kwargs) + ret = f(*args, **kwargs) + for post in posts: + if post.__name__ not in self.ignored_hooks: + post(ret, *args, **kwargs) + return ret + + return run + def nice_output(s): html = f''' @@ -103,26 +141,6 @@ def nice_output(s): return html -def attach_hooks(f, hooks): - pres = [hook_function for (hook_function, hook_type) in hooks if hook_type.lower().startswith('pre')] - posts = [hook_function for (hook_function, hook_type) in hooks if hook_type.lower().startswith('post')] - - @functools.wraps(f) - def run(*args, **kwargs): - caller = traceback.extract_stack()[-2].filename - if caller.startswith(PANDAS_DIR): - ret = f(*args, **kwargs) - else: - for pre in pres: - pre(*args, **kwargs) - ret = f(*args, **kwargs) - for post in posts: - post(ret, *args, **kwargs) - return ret - - return run - - def get_arg(args, kwargs, which_arg, which_kwarg): try: ret = args[which_arg]