mxlpy 0.21.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

mxlpy/meta/source_tools.py

@@ -3,28 +3,166 @@
 from __future__ import annotations
 
 import ast
+import importlib
 import inspect
+import logging
+import math
 import textwrap
 from dataclasses import dataclass
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, cast
 
 import dill
+import numpy as np
 import sympy
-from sympy.printing.pycode import pycode
 
 if TYPE_CHECKING:
     from collections.abc import Callable
-    from types import ModuleType
 
 __all__ = [
     "Context",
+    "KNOWN_CONSTANTS",
+    "KNOWN_FNS",
+    "PARSE_ERROR",
     "fn_to_sympy",
     "get_fn_ast",
     "get_fn_source",
-    "sympy_to_fn",
-    "sympy_to_inline",
 ]
 
+_LOGGER = logging.getLogger(__name__)
+PARSE_ERROR = sympy.Symbol("ERROR")
+
+KNOWN_CONSTANTS: dict[float, sympy.Float] = {
+    math.e: sympy.E,
+    math.pi: sympy.pi,
+    math.nan: sympy.nan,
+    math.tau: sympy.pi * 2,
+    math.inf: sympy.oo,
+    # numpy
+    np.e: sympy.E,
+    np.pi: sympy.pi,
+    np.nan: sympy.nan,
+    np.inf: sympy.oo,
+}
+
+KNOWN_FNS: dict[Callable, sympy.Expr] = {
+    # built-ins
+    abs: sympy.Abs,  # type: ignore
+    min: sympy.Min,
+    max: sympy.Max,
+    pow: sympy.Pow,
+    # round: sympy
+    # divmod
+    # math module
+    math.acos: sympy.acos,
+    math.acosh: sympy.acosh,
+    math.asin: sympy.asin,
+    math.asinh: sympy.asinh,
+    math.atan: sympy.atan,
+    math.atan2: sympy.atan2,
+    math.atanh: sympy.atanh,
+    math.cbrt: sympy.cbrt,
+    math.ceil: sympy.ceiling,
+    # math.comb: sympy.comb,
+    # math.copysign: sympy.copysign,
+    math.cos: sympy.cos,
+    math.cosh: sympy.cosh,
+    # math.degrees: sympy.degrees,
+    # math.dist: sympy.dist,
+    math.erf: sympy.erf,
+    math.erfc: sympy.erfc,
+    math.exp: sympy.exp,
+    # math.exp2: sympy.exp2,
+    # math.expm1: sympy.expm1,
+    # math.fabs: sympy.fabs,
+    math.factorial: sympy.factorial,
+    math.floor: sympy.floor,
+    # math.fmod: sympy.fmod,
+    # math.frexp: sympy.frexp,
+    # math.fsum: sympy.fsum,
+    math.gamma: sympy.gamma,
+    math.gcd: sympy.gcd,
+    # math.hypot: sympy.hypot,
+    # math.isclose: sympy.isclose,
+    # math.isfinite: sympy.isfinite,
+    # math.isinf: sympy.isinf,
+    # math.isnan: sympy.isnan,
+    # math.isqrt: sympy.isqrt,
+    math.lcm: sympy.lcm,
+    # math.ldexp: sympy.ldexp,
+    # math.lgamma: sympy.lgamma,
+    math.log: sympy.log,
+    # math.log10: sympy.log10,
+    # math.log1p: sympy.log1p,
+    # math.log2: sympy.log2,
+    # math.modf: sympy.modf,
+    # math.nextafter: sympy.nextafter,
+    # math.perm: sympy.perm,
+    math.pow: sympy.Pow,
+    math.prod: sympy.prod,
+    math.radians: sympy.rad,
+    math.remainder: sympy.rem,
+    math.sin: sympy.sin,
+    math.sinh: sympy.sinh,
+    math.sqrt: sympy.sqrt,
+    # math.sumprod: sympy.sumprod,
+    math.tan: sympy.tan,
+    math.tanh: sympy.tanh,
+    math.trunc: sympy.trunc,
+    # math.ulp: sympy.ulp,
+    # numpy
+    np.exp: sympy.exp,
+    np.abs: sympy.Abs,
+    np.acos: sympy.acos,
+    np.acosh: sympy.acosh,
+    np.asin: sympy.asin,
+    np.asinh: sympy.asinh,
+    np.atan: sympy.atan,
+    np.atanh: sympy.atanh,
+    np.atan2: sympy.atan2,
+    np.pow: sympy.Pow,
+    np.absolute: sympy.Abs,
+    np.add: sympy.Add,
+    np.arccos: sympy.acos,
+    np.arccosh: sympy.acosh,
+    np.arcsin: sympy.asin,
+    np.arcsinh: sympy.asinh,
+    np.arctan2: sympy.atan2,
+    np.arctan: sympy.atan,
+    np.arctanh: sympy.atanh,
+    np.cbrt: sympy.cbrt,
+    np.ceil: sympy.ceiling,
+    np.conjugate: sympy.conjugate,
+    np.cos: sympy.cos,
+    np.cosh: sympy.cosh,
+    np.exp: sympy.exp,
+    np.floor: sympy.floor,
+    np.gcd: sympy.gcd,
+    np.greater: sympy.GreaterThan,
+    np.greater_equal: sympy.Ge,
+    np.invert: sympy.invert,
+    np.lcm: sympy.lcm,
+    np.less: sympy.LessThan,
+    np.less_equal: sympy.Le,
+    np.log: sympy.log,
+    np.maximum: sympy.maximum,
+    np.minimum: sympy.minimum,
+    np.mod: sympy.Mod,
+    np.positive: sympy.Abs,
+    np.power: sympy.Pow,
+    np.sign: sympy.sign,
+    np.sin: sympy.sin,
+    np.sinh: sympy.sinh,
+    np.sqrt: sympy.sqrt,
+    # np.square: sympy.square,
+    # np.subtract: sympy., # Add(x, -1 * y)
+    np.tan: sympy.tan,
+    np.tanh: sympy.tanh,
+    # np.true_divide: sympy.true_divide,
+    np.trunc: sympy.trunc,
+    # np.vecdot: sympy.vecdot,
+}
+
 
 @dataclass
 class Context:
@@ -33,6 +171,9 @@ class Context:
     symbols: dict[str, sympy.Symbol | sympy.Expr]
     caller: Callable
     parent_module: ModuleType | None
+    origin: str
+    modules: dict[str, ModuleType]
+    fns: dict[str, Callable]
 
     def updated(
         self,
@@ -47,9 +188,23 @@ class Context:
             parent_module=self.parent_module
             if parent_module is None
             else parent_module,
+            origin=self.origin,
+            modules=self.modules,
+            fns=self.fns,
         )
 
 
+def _find_root(value: ast.Attribute | ast.Name, levels: list) -> list[str]:
+    if isinstance(value, ast.Attribute):
+        return _find_root(
+            cast(ast.Attribute, value.value),
+            [value.attr, *levels],
+        )
+
+    root = str(value.id)
+    return [root, *levels]
+
+
 def get_fn_source(fn: Callable) -> str:
     """Get the string representation of a function.
 
@@ -110,121 +265,80 @@ def get_fn_ast(fn: Callable) -> ast.FunctionDef:
     return fn_def
 
 
-def sympy_to_inline(expr: sympy.Expr) -> str:
-    """Convert a sympy expression to inline Python code.
-
-    Parameters
-    ----------
-    expr
-        The sympy expression to convert
-
-    Returns
-    -------
-    str
-        Python code string for the expression
-
-    Examples
-    --------
-    >>> import sympy
-    >>> x = sympy.Symbol('x')
-    >>> expr = x**2 + 2*x + 1
-    >>> sympy_to_inline(expr)
-    'x**2 + 2*x + 1'
-
-    """
-    return cast(str, pycode(expr, fully_qualified_modules=True))
-
-
-def sympy_to_fn(
-    *,
-    fn_name: str,
-    args: list[str],
-    expr: sympy.Expr,
-) -> str:
-    """Convert a sympy expression to a python function.
-
-    Parameters
-    ----------
-    fn_name
-        Name of the function to generate
-    args
-        List of argument names for the function
-    expr
-        Sympy expression to convert to a function body
-
-    Returns
-    -------
-    str
-        String representation of the generated function
-
-    Examples
-    --------
-    >>> import sympy
-    >>> x, y = sympy.symbols('x y')
-    >>> expr = x**2 + y
-    >>> print(sympy_to_fn(fn_name="square_plus_y", args=["x", "y"], expr=expr))
-    def square_plus_y(x: float, y: float) -> float:
-        return x**2 + y
-
-    """
-    fn_args = ", ".join(f"{i}: float" for i in args)
-
-    return f"""def {fn_name}({fn_args}) -> float:
-    return {pycode(expr)}
-    """
-
-
 def fn_to_sympy(
     fn: Callable,
+    origin: str,
     model_args: list[sympy.Symbol | sympy.Expr] | None = None,
-) -> sympy.Expr:
+) -> sympy.Expr | None:
     """Convert a python function to a sympy expression.
 
-    Parameters
-    ----------
-    fn
-        The function to convert
-    model_args
-        Optional list of sympy symbols to substitute for function arguments
+    Args:
+        fn: The function to convert
+        origin: Name of the original caller. Used for error messages.
+        model_args: Optional list of sympy symbols to substitute for function arguments
 
-    Returns
-    -------
-    sympy.Expr
+    Returns:
         Sympy expression equivalent to the function
 
-    Examples
-    --------
-    >>> def square_fn(x):
-    ...     return x**2
-    >>> import sympy
-    >>> fn_to_sympy(square_fn)
-    x**2
-    >>> # With model_args
-    >>> y = sympy.Symbol('y')
-    >>> fn_to_sympy(square_fn, [y])
-    y**2
+    Examples:
+        >>> def square_fn(x):
+        ...     return x**2
+        >>> import sympy
+        >>> fn_to_sympy(square_fn)
+        x**2
+        >>> # With model_args
+        >>> y = sympy.Symbol('y')
+        >>> fn_to_sympy(square_fn, [y])
+        y**2
 
     """
-    fn_def = get_fn_ast(fn)
-    fn_args = [str(arg.arg) for arg in fn_def.args.args]
-    sympy_expr = _handle_fn_body(
-        fn_def.body,
-        ctx=Context(
-            symbols={name: sympy.Symbol(name) for name in fn_args},
-            caller=fn,
-            parent_module=inspect.getmodule(fn),
-        ),
-    )
-    if model_args is not None:
-        sympy_expr = sympy_expr.subs(dict(zip(fn_args, model_args, strict=True)))
-    return cast(sympy.Expr, sympy_expr)
+    try:
+        fn_def = get_fn_ast(fn)
+        fn_args = [str(arg.arg) for arg in fn_def.args.args]
+
+        sympy_expr = _handle_fn_body(
+            fn_def.body,
+            ctx=Context(
+                symbols={name: sympy.Symbol(name) for name in fn_args},
+                caller=fn,
+                parent_module=inspect.getmodule(fn),
+                origin=origin,
+                modules={},
+                fns={},
+            ),
+        )
+        if sympy_expr is None:
+            return None
+        # FIXME: we shouldn't end up here, where does this come from?
+        if isinstance(sympy_expr, float):
+            return sympy.Float(sympy_expr)
+        if model_args is not None:
+            sympy_expr = sympy_expr.subs(dict(zip(fn_args, model_args, strict=True)))
+        return cast(sympy.Expr, sympy_expr)
+
+    except (TypeError, ValueError, NotImplementedError) as e:
+        msg = f"Failed parsing function of {origin}"
+        _LOGGER.warning(msg)
+        _LOGGER.debug("", exc_info=e)
+        return None
 
 
 def _handle_name(node: ast.Name, ctx: Context) -> sympy.Symbol | sympy.Expr:
-    return ctx.symbols[node.id]
+    value = ctx.symbols.get(node.id)
+    if value is None:
+        global_variables = dict(
+            inspect.getmembers(
+                ctx.parent_module,
+                predicate=lambda x: isinstance(x, float),
+            )
+        )
+        value = sympy.Float(global_variables[node.id])
+    return value
 
 
-def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr:
+def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr | None:
+    if isinstance(node, float):
+        return sympy.Float(node)
     if isinstance(node, ast.UnaryOp):
         return _handle_unaryop(node, ctx)
     if isinstance(node, ast.BinOp):
@@ -233,6 +347,11 @@ def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr:
         return _handle_name(node, ctx)
     if isinstance(node, ast.Constant):
         return node.value
+    if isinstance(node, ast.Call):
+        return _handle_call(node, ctx=ctx)
+    if isinstance(node, ast.Attribute):
+        return _handle_attribute(node, ctx=ctx)
+
     if isinstance(node, ast.Compare):
         # Handle chained comparisons like 1 < a < 2
         left = cast(Any, _handle_expr(node.left, ctx))
@@ -263,8 +382,6 @@ def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr:
         for comp in comparisons[1:]:
             result = sympy.And(result, comp)
         return cast(sympy.Expr, result)
-    if isinstance(node, ast.Call):
-        return _handle_call(node, ctx)
 
     # Handle conditional expressions (ternary operators)
     if isinstance(node, ast.IfExp):
@@ -277,7 +394,7 @@ def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr:
     raise NotImplementedError(msg)
 
 
-def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr:
+def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr | None:
     pieces = []
     remaining_body = list(body)
 
@@ -333,7 +450,10 @@ def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr:
                         target_elements, value_elements, strict=True
                     ):
                         if isinstance(target, ast.Name):
-                            ctx.symbols[target.id] = _handle_expr(value_expr, ctx)
+                            expr = _handle_expr(value_expr, ctx)
+                            if expr is None:
+                                return None
+                            ctx.symbols[target.id] = expr
                 else:
                     # Handle potential iterable unpacking
                     value = _handle_expr(node.value, ctx)
@@ -344,8 +464,33 @@ def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr:
                     raise TypeError(msg)
                 target_name = target.id
                 value = _handle_expr(node.value, ctx)
+                if value is None:
+                    return None
                 ctx.symbols[target_name] = value
 
+        elif isinstance(node, ast.Import):
+            for alias in node.names:
+                name = alias.name
+                ctx.modules[name] = importlib.import_module(name)
+
+        elif isinstance(node, ast.ImportFrom):
+            package = cast(str, node.module)
+            module = importlib.import_module(package)
+            contents = dict(inspect.getmembers(module))
+            for alias in node.names:
+                name = alias.name
+                el = contents[name]
+                if isinstance(el, float):
+                    ctx.symbols[name] = sympy.Float(el)
+                elif callable(el):
+                    ctx.fns[name] = el
+                elif isinstance(el, ModuleType):
+                    ctx.modules[name] = el
+                else:
+                    _LOGGER.debug("Skipping import %s", node)
+        else:
+            _LOGGER.debug("Skipping node of type %s", type(node))
+
     # If we have pieces to combine into a Piecewise
     if pieces:
         return sympy.Piecewise(*pieces)
@@ -364,13 +509,14 @@ def _handle_unaryop(node: ast.UnaryOp, ctx: Context) -> sympy.Expr:
     left = _handle_expr(node.operand, ctx)
     left = cast(Any, left)  # stupid sympy types don't allow ops on symbols
 
-    if isinstance(node.op, ast.UAdd):
-        return +left
-    if isinstance(node.op, ast.USub):
-        return -left
-
-    msg = f"Operation {type(node.op).__name__} not implemented"
-    raise NotImplementedError(msg)
+    match node.op:
+        case ast.UAdd():
+            return +left
+        case ast.USub():
+            return -left
+        case _:
+            msg = f"Operation {type(node.op).__name__} not implemented"
+            raise NotImplementedError(msg)
 
 
 def _handle_binop(node: ast.BinOp, ctx: Context) -> sympy.Expr:
@@ -380,63 +526,158 @@ def _handle_binop(node: ast.BinOp, ctx: Context) -> sympy.Expr:
     right = _handle_expr(node.right, ctx)
     right = cast(Any, right)  # stupid sympy types don't allow ops on symbols
 
-    if isinstance(node.op, ast.Add):
-        return left + right
-    if isinstance(node.op, ast.Sub):
-        return left - right
-    if isinstance(node.op, ast.Mult):
-        return left * right
-    if isinstance(node.op, ast.Div):
-        return left / right
-    if isinstance(node.op, ast.Pow):
-        return left**right
-    if isinstance(node.op, ast.Mod):
-        return left % right
-    if isinstance(node.op, ast.FloorDiv):
-        return left // right
-
-    msg = f"Operation {type(node.op).__name__} not implemented"
-    raise NotImplementedError(msg)
+    match node.op:
+        case ast.Add():
+            return left + right
+        case ast.Sub():
+            return left - right
+        case ast.Mult():
+            return left * right
+        case ast.Div():
+            return left / right
+        case ast.Pow():
+            return left**right
+        case ast.Mod():
+            return left % right
+        case ast.FloorDiv():
+            return left // right
+        case _:
+            msg = f"Operation {type(node.op).__name__} not implemented"
+            raise NotImplementedError(msg)
+
+
+# FIXME: check if target isn't an object or class
+def _handle_attribute(node: ast.Attribute, ctx: Context) -> sympy.Expr | None:
+    """Handle an attribute.
+
+    Structures to expect:
+        Attribute(Name(id), attr)             | direct
+        Attribute(Attribute(Name(id)), attr)  | single layer of nesting
+        Attribute(Attribute(...), attr)       | arbitrary nesting
+
+    Targets to expect:
+        - modules (both absolute and relative import)
+            - import a; a.attr
+            - import a; a.b.attr
+            - from a import b; b.attr
+        - objects, e.g. Parameters().a
+        - classes, e.g. Parameters.a
+
+    Watch out for relative imports and the different ways they can be called
+    import a
+    from a import b
+    from a.b import c
+
+    a.attr
+    b.attr
+    c.attr
+    a.b.attr
+    b.c.attr
+    a.b.c.attr
+    """
+    name = str(node.attr)
+    module: ModuleType | None = None
+    modules = (
+        dict(inspect.getmembers(ctx.parent_module, predicate=inspect.ismodule))
+        | ctx.modules
+    )
+    match node.value:
+        case ast.Name(l1):
+            module_name = l1
+            module = modules.get(module_name)
+        case ast.Attribute():
+            levels = _find_root(node.value, [])
+            module_name = ".".join(levels)
+            for level in levels[:-1]:
+                modules.update(
+                    dict(inspect.getmembers(modules[level], predicate=inspect.ismodule))
+                )
+            module = modules.get(levels[-1])
+        case _:
+            raise NotImplementedError
+
+    # Fall-back to absolute import
+    if module is None:
+        module = importlib.import_module(module_name)
+
+    element = dict(
+        inspect.getmembers(
+            module,
+            predicate=lambda x: isinstance(x, float),
+        )
+    ).get(name)
 
+    if element is None:
+        return None
 
-def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr:
-    # direct call, e.g. mass_action(x, k1)
-    if isinstance(callee := node.func, ast.Name):
-        fn_name = str(callee.id)
-        fns = dict(inspect.getmembers(ctx.parent_module, predicate=callable))
+    if (value := KNOWN_CONSTANTS.get(element)) is not None:
+        return value
+    return sympy.Float(element)
 
-        return fn_to_sympy(
-            fns[fn_name],
-            model_args=[_handle_expr(i, ctx) for i in node.args],
-        )
 
-    # search for fn in other namespace
-    if isinstance(attr := node.func, ast.Attribute):
-        imports = dict(inspect.getmembers(ctx.parent_module, inspect.ismodule))
+# FIXME: check if target isn't an object or class
+def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr | None:
+    """Handle call expression.
 
-        # Single level, e.g. fns.mass_action(x, k1)
-        if isinstance(module_name := attr.value, ast.Name):
-            return _handle_call(
-                ast.Call(func=ast.Name(attr.attr), args=node.args, keywords=[]),
-                ctx=ctx.updated(parent_module=imports[module_name.id]),
-            )
+    Variants
+        - mass_action(x, k1)
+        - fns.mass_action(x, k1)
+        - mxlpy.fns.mass_action(x, k1)
 
-        # Multiple levels, e.g. mxlpy.fns.mass_action(x, k1)
-        if isinstance(inner_attr := attr.value, ast.Attribute):
-            if not isinstance(module_name := inner_attr.value, ast.Name):
-                msg = f"Unknown target kind {module_name}"
-                raise NotImplementedError(msg)
-            return _handle_call(
-                ast.Call(
-                    func=ast.Attribute(
-                        value=ast.Name(inner_attr.attr),
-                        attr=attr.attr,
-                    ),
-                    args=node.args,
-                    keywords=[],
-                ),
-                ctx=ctx.updated(parent_module=imports[module_name.id]),
+    In future think about?
+        - object.call
+        - Class.call
+    """
+    model_args: list[sympy.Expr] = []
+    for i in node.args:
+        if (expr := _handle_expr(i, ctx)) is None:
+            return None
+        model_args.append(expr)
+
+    match node.func:
+        case ast.Name(id):
+            fn_name = str(id)
+            fns = (
+                dict(inspect.getmembers(ctx.parent_module, predicate=callable))
+                | ctx.fns
+            )
+            py_fn = fns.get(fn_name)
+
+        # FIXME: use _handle_attribute for this
+        case ast.Attribute(attr=fn_name):
+            module: ModuleType | None = None
+            modules = (
+                dict(inspect.getmembers(ctx.parent_module, predicate=inspect.ismodule))
+                | ctx.modules
             )
 
-    msg = f"Unsupported function type {node.func}"
-    raise NotImplementedError(msg)
+            levels = _find_root(node.func, [])
+            module_name = ".".join(levels[:-1])
+
+            _LOGGER.debug("Searching for module %s", module_name)
+            for level in levels[:-1]:
+                modules.update(
+                    dict(inspect.getmembers(modules[level], predicate=inspect.ismodule))
+                )
+            module = modules.get(levels[-2])
+
+            # Fall-back to absolute import
+            if module is None:
+                module = importlib.import_module(module_name)
+
+            fns = dict(inspect.getmembers(module, predicate=callable))
+            py_fn = fns.get(fn_name)
+        case _:
+            raise NotImplementedError
+
+    if py_fn is None:
+        return None
+
+    if (fn := KNOWN_FNS.get(py_fn)) is not None:
+        return fn
+
+    return fn_to_sympy(
+        py_fn,
+        origin=ctx.origin,
+        model_args=model_args,
+    )