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