ovld 0.3.9__py3-none-any.whl → 0.4.1__py3-none-any.whl

ovld/recode.py

@@ -1,11 +1,293 @@
+import ast
 import inspect
-from types import FunctionType
+import linecache
+import textwrap
+from ast import _splitlines_no_ff as splitlines
+from functools import reduce
+from itertools import count
+from types import CodeType, FunctionType
 
-recurse = object()
+from .dependent import DependentType
+from .utils import Unusable, UsageError
+
+recurse = Unusable(
+    "recurse() can only be used from inside an @ovld-registered function."
+)
+call_next = Unusable(
+    "call_next() can only be used from inside an @ovld-registered function."
+)
+
+
+dispatch_template = """
+from ovld.utils import MISSING
+
+def __DISPATCH__(self, {args}):
+    {inits}
+    {body}
+    {call}
+"""
+
+
+call_template = """
+method = self.map[({lookup})]
+return method({posargs})
+"""
+
+
+def instantiate_code(symbol, code, inject={}):
+    virtual_file = f"<ovld{hash(code)}>"
+    linecache.cache[virtual_file] = (None, None, splitlines(code), virtual_file)
+    code = compile(source=code, filename=virtual_file, mode="exec")
+    glb = {**inject}
+    exec(code, glb, glb)
+    return glb[symbol]
+
+
+# # Previous version: generate a temporary file
+# def instantiate_code(symbol, code, inject={}):
+#     tf = tempfile.NamedTemporaryFile("w")
+#     _tempfiles.append(tf)
+#     tf.write(code)
+#     tf.flush()
+#     glb = runpy.run_path(tf.name)
+#     rval = glb[symbol]
+#     rval.__globals__.update(inject)
+#     return rval
+
+
+def generate_dispatch(arganal):
+    def join(li, sep=", ", trail=False):
+        li = [x for x in li if x]
+        rval = sep.join(li)
+        if len(li) == 1 and trail:
+            rval += ","
+        return rval
+
+    spr, spo, pr, po, kr, ko = arganal.compile()
+
+    inits = set()
+
+    kwargsstar = ""
+    targsstar = ""
+
+    args = []
+    body = [""]
+    posargs = ["self.obj" if arganal.is_method else ""]
+    lookup = []
+
+    i = 0
+
+    for name in spr + spo:
+        if name in spr:
+            args.append(name)
+        else:
+            args.append(f"{name}=MISSING")
+        posargs.append(name)
+        lookup.append(f"{arganal.lookup_for(i)}({name})")
+        i += 1
+
+    if len(po) <= 1:
+        # If there are more than one non-strictly positional optional arguments,
+        # then all positional arguments are strictly positional, because if e.g.
+        # x and y are optional we want x==MISSING to imply that y==MISSING, but
+        # that only works if y cannot be provided as a keyword argument.
+        args.append("/")
+
+    for name in pr + po:
+        if name in pr:
+            args.append(name)
+        else:
+            args.append(f"{name}=MISSING")
+        posargs.append(name)
+        lookup.append(f"{arganal.lookup_for(i)}({name})")
+        i += 1
+
+    if len(po) > 1:
+        args.append("/")
+
+    if kr or ko:
+        args.append("*")
+
+    for name in kr:
+        lookup_fn = (
+            "self.map.transform"
+            if name in arganal.complex_transforms
+            else "type"
+        )
+        args.append(f"{name}")
+        posargs.append(f"{name}={name}")
+        lookup.append(f"({name!r}, {lookup_fn}({name}))")
+
+    for name in ko:
+        args.append(f"{name}=MISSING")
+        kwargsstar = "**KWARGS"
+        targsstar = "*TARGS"
+        inits.add("KWARGS = {}")
+        inits.add("TARGS = []")
+        body.append(f"if {name} is not MISSING:")
+        body.append(f"    KWARGS[{name!r}] = {name}")
+        body.append(
+            f"    TARGS.append(({name!r}, {arganal.lookup_for(name)}({name})))"
+        )
+
+    posargs.append(kwargsstar)
+    lookup.append(targsstar)
+
+    fullcall = call_template.format(
+        lookup=join(lookup, trail=True),
+        posargs=join(posargs),
+    )
+
+    calls = []
+    if spo or po:
+        req = len(spr + pr)
+        for i, arg in enumerate(spo + po):
+            call = call_template.format(
+                lookup=join(lookup[: req + i], trail=True),
+                posargs=join(posargs[: req + i + 1]),
+            )
+            call = textwrap.indent(call, "    ")
+            calls.append(f"\nif {arg} is MISSING:{call}")
+    calls.append(fullcall)
+
+    code = dispatch_template.format(
+        inits=join(inits, sep="\n    "),
+        args=join(args),
+        body=join(body, sep="\n    "),
+        call=textwrap.indent("".join(calls), "    "),
+    )
+    return instantiate_code("__DISPATCH__", code)
+
+
+class GenSym:
+    def __init__(self, prefix):
+        self.prefix = prefix
+        self.count = count()
+        self.variables = {}
+
+    def add(self, value):
+        if isinstance(value, (int, float, str)):
+            return repr(value)
+        id = f"{self.prefix}{next(self.count)}"
+        self.variables[id] = value
+        return id
+
+
+def generate_dependent_dispatch(tup, handlers, next_call, slf, name, err, nerr):
+    def to_dict(tup):
+        return dict(
+            entry if isinstance(entry, tuple) else (i, entry)
+            for i, entry in enumerate(tup)
+        )
+
+    def argname(x):
+        return f"ARG{x}" if isinstance(x, int) else x
+
+    def argprovide(x):
+        return f"ARG{x}" if isinstance(x, int) else f"{x}={x}"
+
+    def codegen(typ, arg):
+        cg = typ.codegen()
+        return cg.template.format(
+            arg=arg, **{k: gen.add(v) for k, v in cg.substitutions.items()}
+        )
+
+    tup = to_dict(tup)
+    handlers = [(h, to_dict(types)) for h, types in handlers]
+    gen = GenSym(prefix="INJECT")
+    conjs = []
+
+    exclusive = False
+    keyexpr = None
+    keyed = None
+    for k in tup:
+        featured = set(types[k] for h, types in handlers)
+        if len(featured) == len(handlers):
+            possibilities = set(type(t) for t in featured)
+            focus = possibilities.pop()
+            # Possibilities is now empty if only one type of DependentType
+
+            if not possibilities:
+                if getattr(focus, "keyable_type", False):
+                    all_keys = [
+                        {key: h for key in types[k].get_keys()}
+                        for h, types in handlers
+                    ]
+                    keyed = reduce(lambda a, b: {**a, **b}, all_keys)
+                    if (
+                        len(keyed) == sum(map(len, all_keys))
+                        and len(featured) < 4
+                    ):
+                        exclusive = True
+                        keyexpr = None
+                    else:
+                        keyexpr = focus.keygen().format(arg=argname(k))
+
+                else:
+                    exclusive = getattr(focus, "exclusive_type", False)
+
+    for i, (h, types) in enumerate(handlers):
+        relevant = [k for k in tup if isinstance(types[k], DependentType)]
+        if len(relevant) > 1:
+            # The keyexpr method only works if there is only one condition to check.
+            keyexpr = keyed = None
+        codes = [codegen(types[k], argname(k)) for k in relevant]
+        conj = " and ".join(codes)
+        if not conj:  # pragma: no cover
+            # Not sure if this can happen
+            conj = "True"
+        conjs.append(conj)
+
+    argspec = ", ".join(argname(x) for x in tup)
+    argcall = ", ".join(argprovide(x) for x in tup)
+
+    body = []
+    if keyexpr:
+        body.append(f"HANDLER = {gen.add(keyed)}.get({keyexpr}, FALLTHROUGH)")
+        body.append(f"return HANDLER({slf}{argcall})")
+
+    elif exclusive:
+        for i, conj in enumerate(conjs):
+            body.append(f"if {conj}: return HANDLER{i}({slf}{argcall})")
+        body.append(f"return FALLTHROUGH({slf}{argcall})")
+
+    else:
+        for i, conj in enumerate(conjs):
+            body.append(f"MATCH{i} = {conj}")
+
+        summation = " + ".join(f"MATCH{i}" for i in range(len(handlers)))
+        body.append(f"SUMMATION = {summation}")
+        body.append("if SUMMATION == 1:")
+        for i, (h, types) in enumerate(handlers):
+            body.append(f"    if MATCH{i}: return HANDLER{i}({slf}{argcall})")
+        body.append("elif SUMMATION == 0:")
+        body.append(f"    return FALLTHROUGH({slf}{argcall})")
+        body.append("else:")
+        body.append(f"    raise {gen.add(err)}")
+
+    body_text = textwrap.indent("\n".join(body), "    ")
+    code = f"def __DEPENDENT_DISPATCH__({slf}{argspec}):\n{body_text}"
+
+    inject = gen.variables
+    for i, (h, types) in enumerate(handlers):
+        inject[f"HANDLER{i}"] = h
+
+    def raise_error(*args, **kwargs):
+        raise nerr
+
+    inject["FALLTHROUGH"] = (next_call and next_call[0]) or raise_error
+
+    fn = instantiate_code(
+        symbol="__DEPENDENT_DISPATCH__", code=code, inject=inject
+    )
+    return rename_function(fn, name)
+
+
+_current = count()
 
 
 class Conformer:
-    __slots__ = ("code", "orig_fn", "renamed_fn", "ovld", "code2")
+    __slots__ = ("code", "orig_fn", "renamed_fn", "ovld")
 
     def __init__(self, ovld, orig_fn, renamed_fn):
         self.ovld = ovld
@@ -21,21 +303,22 @@ class Conformer:
             new_fn = None
             new_code = new
 
-        if new_code is None:
-            self.ovld.unregister(self.orig_fn)
-
-        elif new_fn is None:  # pragma: no cover
-            # Not entirely sure if this ever happens
-            self.renamed_fn.__code__ = new_code
-
-        elif inspect.signature(self.orig_fn) != inspect.signature(new_fn):
-            self.ovld.unregister(self.orig_fn)
-            self.ovld.register(new_fn)
+        self.ovld.unregister(self.orig_fn)
 
-        else:
-            self.renamed_fn.__code__ = rename_code(
-                new_code, self.renamed_fn.__code__.co_name
+        if new_fn is None:  # pragma: no cover
+            if new_code is None:
+                return
+            ofn = self.orig_fn
+            new_fn = FunctionType(
+                new_code,
+                ofn.__globals__,
+                ofn.__name__,
+                ofn.__defaults__,
+                ofn.__closure__,
             )
+            new_fn.__annotations__ = ofn.__annotations__
+
+        self.ovld.register(new_fn)
 
         from codefind import code_registry
 
@@ -76,5 +359,224 @@ def rename_function(fn, newname):
     new_fn = FunctionType(
         newcode, fn.__globals__, newname, fn.__defaults__, fn.__closure__
     )
+    new_fn.__kwdefaults__ = fn.__kwdefaults__
+    new_fn.__annotations__ = fn.__annotations__
+    return new_fn
+
+
+class NameConverter(ast.NodeTransformer):
+    def __init__(
+        self, anal, recurse_sym, call_next_sym, ovld_mangled, code_mangled
+    ):
+        self.analysis = anal
+        self.recurse_sym = recurse_sym
+        self.call_next_sym = call_next_sym
+        self.ovld_mangled = ovld_mangled
+        self.code_mangled = code_mangled
+
+    def visit_Name(self, node):
+        if node.id == self.recurse_sym:
+            return ast.copy_location(
+                old_node=node,
+                new_node=ast.Name(self.ovld_mangled, ctx=node.ctx),
+            )
+        elif node.id == self.call_next_sym:
+            raise UsageError("call_next should be called right away")
+        else:
+            return node
+
+    def visit_Call(self, node):
+        if not isinstance(node.func, ast.Name) or node.func.id not in (
+            self.recurse_sym,
+            self.call_next_sym,
+        ):
+            return self.generic_visit(node)
+
+        if any(isinstance(arg, ast.Starred) for arg in node.args):
+            return self.generic_visit(node)
+
+        cn = node.func.id == self.call_next_sym
+
+        def _make_lookup_call(key, arg):
+            value = ast.NamedExpr(
+                target=ast.Name(id=f"__TMP{key}", ctx=ast.Store()),
+                value=self.visit(arg),
+            )
+            if self.analysis.lookup_for(key) == "self.map.transform":
+                func = ast.Attribute(
+                    value=ast.Name(id="__TMPM", ctx=ast.Load()),
+                    attr="transform",
+                    ctx=ast.Load(),
+                )
+            else:
+                func = ast.Name(id="type", ctx=ast.Load())
+            return ast.Call(
+                func=func,
+                args=[value],
+                keywords=[],
+            )
+
+        # type index for positional arguments
+        type_parts = [
+            _make_lookup_call(i, arg) for i, arg in enumerate(node.args)
+        ]
+
+        # type index for keyword arguments
+        type_parts += [
+            ast.Tuple(
+                elts=[
+                    ast.Constant(value=kw.arg),
+                    _make_lookup_call(kw.arg, kw.value),
+                ],
+                ctx=ast.Load(),
+            )
+            for kw in node.keywords
+        ]
+
+        if cn:
+            type_parts.insert(0, ast.Name(id=self.code_mangled, ctx=ast.Load()))
+        method = ast.Subscript(
+            value=ast.NamedExpr(
+                target=ast.Name(id="__TMPM", ctx=ast.Store()),
+                value=ast.Attribute(
+                    value=ast.Name(id=self.ovld_mangled, ctx=ast.Load()),
+                    attr="map",
+                    ctx=ast.Load(),
+                ),
+            ),
+            slice=ast.Tuple(
+                elts=type_parts,
+                ctx=ast.Load(),
+            ),
+            ctx=ast.Load(),
+        )
+        if self.analysis.is_method:
+            method = ast.Call(
+                func=ast.Attribute(
+                    value=method,
+                    attr="__get__",
+                    ctx=ast.Load(),
+                ),
+                args=[ast.Name(id="self", ctx=ast.Load())],
+                keywords=[],
+            )
+
+        new_node = ast.Call(
+            func=method,
+            args=[
+                ast.Name(id=f"__TMP{i}", ctx=ast.Load())
+                for i, arg in enumerate(node.args)
+            ],
+            keywords=[
+                ast.keyword(
+                    arg=kw.arg,
+                    value=ast.Name(id=f"__TMP{kw.arg}", ctx=ast.Load()),
+                )
+                for kw in node.keywords
+            ],
+        )
+        return ast.copy_location(old_node=node, new_node=new_node)
+
+
+def _search_names(co, values, glb, closure=None):
+    if isinstance(co, CodeType):
+        if closure is not None:
+            for varname, cell in zip(co.co_freevars, closure):
+                if any(cell.cell_contents is v for v in values):
+                    yield varname
+        for name in co.co_names:
+            if any(glb.get(name, None) is v for v in values):
+                yield name
+        else:
+            for ct in co.co_consts:
+                yield from _search_names(ct, values, glb)
+
+
+def adapt_function(fn, ovld, newname):
+    """Create a copy of the function with a different name."""
+    rec_syms = list(
+        _search_names(
+            fn.__code__, (recurse, ovld), fn.__globals__, fn.__closure__
+        )
+    )
+    cn_syms = list(
+        _search_names(fn.__code__, (call_next,), fn.__globals__, fn.__closure__)
+    )
+    if rec_syms or cn_syms:
+        return recode(
+            fn, ovld, rec_syms and rec_syms[0], cn_syms and cn_syms[0], newname
+        )
+    else:
+        return rename_function(fn, newname)
+
+
+def closure_wrap(tree, fname, names):
+    wrap = ast.copy_location(
+        ast.FunctionDef(
+            name="##create_closure",
+            args=ast.arguments(
+                posonlyargs=[],
+                args=[ast.arg(arg=name) for name in names],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=[],
+                kwarg=None,
+                defaults=[],
+            ),
+            body=[
+                tree,
+                ast.Return(ast.Name(id=fname, ctx=ast.Load())),
+            ],
+            decorator_list=[],
+            returns=None,
+        ),
+        tree,
+    )
+    ast.fix_missing_locations(wrap)
+    return ast.Module(body=[wrap], type_ignores=[])
+
+
+def recode(fn, ovld, recurse_sym, call_next_sym, newname):
+    ovld_mangled = f"___OVLD{ovld.id}"
+    code_mangled = f"___CODE{next(_current)}"
+    try:
+        src = inspect.getsource(fn)
+    except OSError:  # pragma: no cover
+        raise OSError(
+            f"ovld is unable to rewrite {fn} because it cannot read its source code."
+            " It may be an issue with __pycache__, so try to either change the source"
+            " to force a refresh, or remove __pycache__ altogether. If that does not work,"
+            " avoid calling recurse()/call_next()"
+        )
+    tree = ast.parse(textwrap.dedent(src))
+    new = NameConverter(
+        anal=ovld.argument_analysis,
+        recurse_sym=recurse_sym,
+        call_next_sym=call_next_sym,
+        ovld_mangled=ovld_mangled,
+        code_mangled=code_mangled,
+    ).visit(tree)
+    new.body[0].decorator_list = []
+    if fn.__closure__:
+        new = closure_wrap(new.body[0], "irrelevant", fn.__code__.co_freevars)
+    ast.fix_missing_locations(new)
+    ast.increment_lineno(new, fn.__code__.co_firstlineno - 1)
+    res = compile(new, mode="exec", filename=fn.__code__.co_filename)
+    if fn.__closure__:
+        res = [x for x in res.co_consts if isinstance(x, CodeType)][0]
+    (*_, new_code) = [ct for ct in res.co_consts if isinstance(ct, CodeType)]
+    new_closure = tuple(
+        [
+            fn.__closure__[fn.__code__.co_freevars.index(name)]
+            for name in new_code.co_freevars
+        ]
+    )
+    new_fn = FunctionType(
+        new_code, fn.__globals__, newname, fn.__defaults__, new_closure
+    )
+    new_fn.__kwdefaults__ = fn.__kwdefaults__
     new_fn.__annotations__ = fn.__annotations__
+    new_fn = rename_function(new_fn, newname)
+    new_fn.__globals__[ovld_mangled] = ovld
+    new_fn.__globals__[code_mangled] = new_fn.__code__
     return new_fn