Source code for dataclass_wizard.utils.function_builder

from dataclasses import MISSING

from ..log import LOG


[docs] def is_builtin_class(cls: type) -> bool: """Check if a class is a builtin in Python.""" return cls.__module__ == 'builtins'
[docs] class FunctionBuilder: __slots__ = ( 'current_function', 'prev_function', 'functions', 'globals', 'indent_level', 'namespace', ) def __init__(self): self.functions = {} self.indent_level = 0 self.globals = {} self.namespace = {} def __ior__(self, other): """ Allows `|=` operation for :class:`FunctionBuilder` objects, e.g. :: my_fn_builder |= other_fn_builder """ self.functions |= other.functions return self def __enter__(self): self.indent_level += 1 def __exit__(self, exc_type, exc_val, exc_tb): indent_lvl = self.indent_level = self.indent_level - 1 if not indent_lvl: self.finalize_function() # noinspection PyAttributeOutsideInit
[docs] def function(self, name: str, args: list, return_type=MISSING, locals=None) -> 'FunctionBuilder': """Start a new function definition with optional return type.""" curr_fn = getattr(self, 'current_function', None) if curr_fn is not None: curr_fn['indent_level'] = self.indent_level self.prev_function = curr_fn self.current_function = { "name": name, "args": args, "body": [], "return_type": return_type, "locals": locals if locals is not None else {}, } self.indent_level = 0 return self
def _with_new_block(self, name: str, condition: 'str | None' = None, comment: str = '') -> 'FunctionBuilder': """Creates a new block. Used with a context manager (with).""" indent = ' ' * self.indent_level if comment: comment = f' # {comment}' if condition is not None: self.current_function["body"].append(f"{indent}{name} {condition}:{comment}") else: self.current_function["body"].append(f"{indent}{name}:{comment}") return self
[docs] def for_(self, condition: str) -> 'FunctionBuilder': """Equivalent to the `for` statement in Python. Sample Usage: >>> with FunctionBuilder().for_('i in range(3)'): >>> ... Will generate the following code: >>> for i in range(3): >>> ... """ return self._with_new_block('for', condition)
[docs] def if_(self, condition: str, comment: str = '') -> 'FunctionBuilder': """Equivalent to the `if` statement in Python. Sample Usage: >>> with FunctionBuilder().if_('something is True'): >>> ... Will generate the following code: >>> if something is True: >>> ... """ return self._with_new_block('if', condition, comment)
[docs] def elif_(self, condition: str) -> 'FunctionBuilder': """Equivalent to the `elif` statement in Python. Sample Usage: >>> with FunctionBuilder().elif_('something is True'): >>> ... Will generate the following code: >>> elif something is True: >>> ... """ return self._with_new_block('elif', condition)
[docs] def else_(self) -> 'FunctionBuilder': """Equivalent to the `else` statement in Python. Sample Usage: >>> with FunctionBuilder().else_(): >>> ... Will generate the following code: >>> else: >>> ... """ return self._with_new_block('else')
[docs] def try_(self) -> 'FunctionBuilder': """Equivalent to the `try` block in Python. Sample Usage: >>> with FunctionBuilder().try_(): >>> ... Will generate the following code: >>> try: >>> ... """ return self._with_new_block('try')
[docs] def except_(self, cls: type[Exception], var_name: 'str | None' = None, *custom_classes: type[Exception]): """Equivalent to the `except` block in Python. Sample Usage: >>> with FunctionBuilder().except_(TypeError, 'exc'): >>> ... Will generate the following code: >>> except TypeError as exc: >>> ... """ cls_name = cls.__name__ statement = f'{cls_name} as {var_name}' if var_name else cls_name if not is_builtin_class(cls): if cls_name not in self.globals: # TODO # LOG.debug('Ensuring class in globals, cls=%s', cls_name) self.globals[cls_name] = cls if custom_classes: for cls in custom_classes: if not is_builtin_class(cls): cls_name = cls.__name__ if cls_name not in self.globals: # LOG.debug('Ensuring class in globals, cls=%s', cls_name) self.globals[cls_name] = cls return self._with_new_block('except', statement)
[docs] def except_multi(self, *classes: type[Exception]): """Equivalent to the `except` block in Python. Sample Usage: >>> with FunctionBuilder().except_multi(AttributeError, TypeError, ValueError): >>> ... Will generate the following code: >>> except (AttributeError, TypeError, ValueError): >>> ... """ if len(classes) == 1: statement = classes[0].__name__ else: class_names = ', '.join([cls.__name__ for cls in classes]) statement = f'({class_names})' return self._with_new_block('except', statement)
[docs] def break_(self): """Equivalent to the `break` statement in Python.""" self.add_line('break')
[docs] def add_line(self, line: str): """Add a line to the current function's body with proper indentation.""" indent = ' ' * self.indent_level self.current_function["body"].append(f"{indent}{line}")
[docs] def add_lines(self, *lines: str): """Add lines to the current function's body with proper indentation.""" indent = ' ' * self.indent_level self.current_function["body"].extend( [f"{indent}{line}" for line in lines] )
[docs] def increase_indent(self): # pragma: no cover """Increase indentation level for nested code.""" self.indent_level += 1
[docs] def decrease_indent(self): # pragma: no cover """Decrease indentation level.""" if self.indent_level > 1: self.indent_level -= 1
[docs] def finalize_function(self): """Finalize the function code and add to the list of functions.""" # Add the function body and don't re-add the function definition curr_fn = self.current_function func_code = '\n'.join(curr_fn["body"]) self.functions[curr_fn["name"]] = { "args": curr_fn["args"], "return_type": curr_fn["return_type"], "locals": curr_fn["locals"], "code": func_code } if (prev_fn := getattr(self, 'prev_function', None)) is not None: self.indent_level = prev_fn.pop('indent_level') self.current_function = prev_fn self.prev_function = None else: self.current_function # Reset current function
[docs] def create_functions(self, _globals=None): """Create functions by compiling the code.""" # Note that we may mutate locals. Callers beware! # The only callers are internal to this module, so no # worries about external callers. # Compute the text of the entire function. # txt = f' def {name}({args}){return_annotation}:\n{body}' # Build the function code for all functions # Free variables in exec are resolved in the global namespace. # The global namespace we have is user-provided, so we can't modify it for # our purposes. So we put the things we need into locals and introduce a # scope to allow the function we're creating to close over them. fn_name_locals_and_code = [] for name, func in self.functions.items(): args = ','.join(func['args']) body = func['code'] return_type = func['return_type'] locals = func['locals'] return_annotation = '' if return_type is not MISSING: locals[f'__dataclass_{name}_return_type__'] = return_type return_annotation = f'->__dataclass_{name}_return_type__' fn_name_locals_and_code.append( (name, locals, f'def {name}({args}){return_annotation}:\n{body}') ) txt = '\n'.join([ f"def __create_{name}_fn__({', '.join(locals.keys())}):\n" f" {code}\n" f" return {name}" for name, locals, code in fn_name_locals_and_code ]) # Print the generated code for debugging # logging.debug(f"Generated function code:\n{all_func_code}") LOG.debug("Generated function code:\n%s", txt) ns = {} # TODO _globals = self.globals if _globals is None else _globals | self.globals LOG.debug("Globals before function compilation: %s", _globals) exec(txt, _globals, ns) # TODO do we need self.namespace? final_ns = self.namespace = {} # TODO: add function to dependent function `locals` rather than to `globals` for name, locals, _ in fn_name_locals_and_code: _globals[name] = final_ns[name] = ns[f'__create_{name}_fn__'](**locals) # final_ns = self.namespace = { # name: ns[f'__create_{name}_fn__'](**locals) # for name, locals, _ in fn_name_locals_and_code # } # Print namespace for debugging LOG.debug("Namespace after function compilation: %s", final_ns) return final_ns