diff --git a/mathics/core/util.py b/mathics/core/util.py index e8c82ea3d..bad2cbe00 100644 --- a/mathics/core/util.py +++ b/mathics/core/util.py @@ -3,10 +3,12 @@ Miscellaneous mathics.core utility functions. """ +import os import sys from itertools import chain from pathlib import PureWindowsPath from platform import python_implementation +import time from typing import Optional IS_PYPY = python_implementation() == "PyPy" @@ -113,3 +115,80 @@ def subranges( items[start : start + length], (items[:start], items[start + length :]), ) + + +def prt_expr_tree(expr, indent=1): + """ + Print a Mathics Expression as an indented tree + """ + if not hasattr(expr, "elements"): + print(" " * indent + str(expr)) + else: + print(" " * indent + str(expr.head)) + for elt in expr.elements: + prt_expr_tree(elt, indent + 1) + + +def prt_sympy_tree(expr, indent=""): + """ + Print a SymPy expression as an indented tree + """ + if expr.args: + print(f"{indent}{expr.func.__name__}") + for i, arg in enumerate(expr.args): + prt_sympy_tree(arg, indent + " ") + else: + print(f"{indent}{expr.func.__name__}({str(expr)})") + + +class Timer: + """ + Times a block of code. May be be used as a decorator or as a context manager: + + # decorator + @Timer(name): + def f(...): + ... + + # context manager + with Timer(name): + ... + + Timings are nested (in execution order), and the output prints the nested + timings as an "upside-down" indented outline, with an outer level printed after + all nested inner levels, supporting both detailed and summary timings. + + Timing.level controls how deeply nested timings are displayed: + -1 all, 0 none, 1 only top level, etc. Default is 0. Use MATHICS_TIMING + environment variable to change. + """ + + level = int(os.getenv("MATHICS_TIMING", "0")) + timers = [] + + def __init__(self, name): + self.name = name + + def __call__(self, fun): + def timed_fun(*args, **kwargs): + with self: + return fun(*args, **kwargs) + return timed_fun + + def start(name): + Timer.timers.append((name, time.time())) + + def stop(): + name, start = Timer.timers.pop() + ms = (time.time() - start) * 1000 + if Timer.level < 0 or len(Timer.timers) < Timer.level: + print(f"{' '*len(Timer.timers)}{name}: {ms:.1f} ms") + + def __enter__(self): + if self.name: + Timer.start(self.name) + + def __exit__(self, *args): + if self.name: + Timer.stop() +