import contextlib import io import re def extract_function_name(text): """Find a function definition and return its name. text: String returns: String or None """ pattern = r"def\s+(\w+)\s*\(" match = re.search(pattern, text) if match: func_name = match.group(1) return func_name else: return None # the functions that define cell magic commands are only defined # if we're running in Jupyter. try: from IPython.core.magic import register_cell_magic from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring @register_cell_magic def add_method_to(args, cell): # get the name of the function defined in this cell func_name = extract_function_name(cell) if func_name is None: return f"This cell doesn't define any new functions." # get the class we're adding it to namespace = get_ipython().user_ns class_name = args cls = namespace.get(class_name, None) if cls is None: return f"Class '{class_name}' not found." # save the old version of the function if it was already defined old_func = namespace.get(func_name, None) if old_func is not None: del namespace[func_name] # Execute the cell to define the function get_ipython().run_cell(cell) # get the newly defined function new_func = namespace.get(func_name, None) if new_func is None: return f"This cell didn't define {func_name}." # add the function to the class and remove it from the namespace setattr(cls, func_name, new_func) del namespace[func_name] # restore the old function to the namespace if old_func is not None: namespace[func_name] = old_func @register_cell_magic def expect_error(line, cell): try: get_ipython().run_cell(cell) except Exception as e: get_ipython().run_cell("%tb") @magic_arguments() @argument("exception", help="Type of exception to catch") @register_cell_magic def expect(line, cell): args = parse_argstring(expect, line) exception = eval(args.exception) try: get_ipython().run_cell(cell) except exception as e: get_ipython().run_cell("%tb") def traceback(mode): """Set the traceback mode. mode: string """ with contextlib.redirect_stdout(io.StringIO()): get_ipython().run_cell(f"%xmode {mode}") traceback("Minimal") except (ImportError, NameError): print("Warning: IPython is not available, cell magic not defined.")