diff --git a/src/calibre/db/tests/reading.py b/src/calibre/db/tests/reading.py index a57dfbc740..040bdeba3d 100644 --- a/src/calibre/db/tests/reading.py +++ b/src/calibre/db/tests/reading.py @@ -847,6 +847,32 @@ def evaluate(book, ctx): v = formatter.safe_format(template, {}, 'TEMPLATE ERROR', mi) self.assertEqual(set(v.split(',')), {'Tag One', 'News', 'Tag Two'}) + # test using a custom context class + template = '''python: +def evaluate(book, ctx): + tags = ctx.db.new_api.all_field_names('tags') + return ','.join(list(ctx.helper_function(tags))) +''' + from calibre.utils.formatter import PythonTemplateContext + class CustomContext(PythonTemplateContext): + def helper_function(self, arg): + s = set(arg) + s.add('helper called') + return s + + v = formatter.safe_format(template, {}, 'TEMPLATE ERROR', mi, + python_context_object=CustomContext()) + self.assertEqual(set(v.split(',')), {'Tag One', 'News', 'Tag Two','helper called'}) + + # test is_multiple values + template = '''python: +def evaluate(book, ctx): + tags = ctx.db.new_api.all_field_names('tags') + return ','.join(list(tags)) +''' + v = formatter.safe_format(template, {}, 'TEMPLATE ERROR', mi) + self.assertEqual(set(v.split(',')), {'Tag One', 'News', 'Tag Two'}) + # test calling a python stored template from a GPM template from calibre.utils.formatter_functions import ( load_user_template_functions, unload_user_template_functions) diff --git a/src/calibre/utils/formatter.py b/src/calibre/utils/formatter.py index 0d46a16889..51eda846bc 100644 --- a/src/calibre/utils/formatter.py +++ b/src/calibre/utils/formatter.py @@ -832,23 +832,23 @@ class StopException(Exception): class PythonTemplateContext(object): - def __init__(self, **kwargs): + def __init__(self): # Set attributes we already know must exist. + object.__init__(self) self.db = None self.arguments = None self.globals = None + self.attrs_set = {'db', 'arguments', 'globals'} - attrs_set = {'db', 'arguments', 'globals'} - + def set_values(self, **kwargs): # Create/set attributes from the named parameters. Doing it this way we - # aren't required to change the signature of __init__ if/when we add + # aren't required to change the signature of this method if/when we add # attributes in the future. However, if a user depends upon the # existence of some attribute and the context creator doesn't supply it # then the user will get an AttributeError exception. for k,v in kwargs.items(): - attrs_set.add(k) + self.attrs_set.add(k) setattr(self, k, v) - self.attrs_set = attrs_set @property def attributes(self): @@ -1598,11 +1598,11 @@ class TemplateFormatter(string.Formatter): def _run_python_template(self, compiled_template, arguments): try: - return compiled_template(self.book, - PythonTemplateContext( - db=get_database(self.book, get_database(self.book, None)), - globals=self.global_vars, - arguments=arguments)) + self.python_context_object.set_values( + db=get_database(self.book, get_database(self.book, None)), + globals=self.global_vars, + arguments=arguments) + return compiled_template(self.book, self.python_context_object) except Exception as e: ss = traceback.extract_tb(exc_info()[2])[-1] raise ValueError(_('Error in function {0} on line {1} : {2} - {3}').format( @@ -1776,7 +1776,8 @@ class TemplateFormatter(string.Formatter): # ######### a formatter that throws exceptions ############ - def unsafe_format(self, fmt, kwargs, book, strip_results=True, global_vars=None): + def unsafe_format(self, fmt, kwargs, book, strip_results=True, global_vars=None, + python_context_object=None): state = self.save_state() try: self.strip_results = strip_results @@ -1786,6 +1787,10 @@ class TemplateFormatter(string.Formatter): self.composite_values = {} self.locals = {} self.global_vars = global_vars if isinstance(global_vars, dict) else {} + if isinstance(python_context_object, PythonTemplateContext): + self.python_context_object = python_context_object + else: + self.python_context_object = PythonTemplateContext() return self.evaluate(fmt, [], kwargs, self.global_vars) finally: self.restore_state(state) @@ -1795,7 +1800,8 @@ class TemplateFormatter(string.Formatter): def safe_format(self, fmt, kwargs, error_value, book, column_name=None, template_cache=None, strip_results=True, template_functions=None, - global_vars=None, break_reporter=None): + global_vars=None, break_reporter=None, + python_context_object=None): state = self.save_state() if self.recursion_level == 0: # Initialize the composite values dict if this is the base-level @@ -1808,6 +1814,10 @@ class TemplateFormatter(string.Formatter): self.kwargs = kwargs self.book = book self.global_vars = global_vars if isinstance(global_vars, dict) else {} + if isinstance(python_context_object, PythonTemplateContext): + self.python_context_object = python_context_object + else: + self.python_context_object = PythonTemplateContext() if template_functions: self.funcs = template_functions else: