ovld 0.4.5__py3-none-any.whl → 0.5.0__py3-none-any.whl

ovld/recode.py

@@ -1,41 +1,23 @@
 import ast
 import inspect
-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
 
-from .utils import MISSING, NameDatabase, Unusable, UsageError, subtler_type
-
-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."
+from .codegen import (
+    Code,
+    instantiate_code,
+    rename_code,
+    rename_function,
+    transfer_function,
 )
+from .utils import MISSING, NameDatabase, SpecialForm, UsageError, subtler_type
 
-
-def instantiate_code(symbol, code, inject={}):
-    virtual_file = f"<ovld:{abs(hash(code)):x}>"
-    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
+recurse = SpecialForm("recurse")
+call_next = SpecialForm("call_next")
+resolve = SpecialForm("resolve")
+current_code = SpecialForm("current_code")
 
 
 dispatch_template = """
@@ -53,6 +35,13 @@ return {mvar}({posargs})
 """
 
 
+def generate_checking_code(typ):
+    if hasattr(typ, "codegen"):
+        return typ.codegen()
+    else:
+        return Code("isinstance($arg, $this)", this=typ)
+
+
 def generate_dispatch(ov, arganal):
     def join(li, sep=", ", trail=False):
         li = [x for x in li if x]
@@ -171,12 +160,11 @@ def generate_dispatch(ov, arganal):
 
 
 def generate_dependent_dispatch(tup, handlers, next_call, slf, name, err, nerr):
-    from .dependent import generate_checking_code, is_dependent
+    from .dependent import is_dependent
 
     def to_dict(tup):
         return dict(
-            entry if isinstance(entry, tuple) else (i, entry)
-            for i, entry in enumerate(tup)
+            entry if isinstance(entry, tuple) else (i, entry) for i, entry in enumerate(tup)
         )
 
     def argname(x):
@@ -185,12 +173,6 @@ def generate_dependent_dispatch(tup, handlers, next_call, slf, name, err, nerr):
     def argprovide(x):
         return f"ARG{x}" if isinstance(x, int) else f"{x}={x}"
 
-    def codegen(typ, arg):
-        cg = generate_checking_code(typ)
-        return cg.template.format(
-            arg=arg, **{k: ndb[v] for k, v in cg.substitutions.items()}
-        )
-
     tup = to_dict(tup)
     handlers = [(h, to_dict(types)) for h, types in handlers]
     ndb = NameDatabase(default_name="INJECT")
@@ -209,18 +191,14 @@ def generate_dependent_dispatch(tup, handlers, next_call, slf, name, err, nerr):
             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
+                        {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
-                    ):
+                    if len(keyed) == sum(map(len, all_keys)) and len(featured) < 4:
                         exclusive = True
                         keyexpr = None
                     else:
-                        keyexpr = focus.keygen().format(arg=argname(k))
+                        keyexpr = focus.keygen().sub(arg=Code(argname(k))).fill(ndb)
 
                 else:
                     exclusive = getattr(focus, "exclusive_type", False)
@@ -230,7 +208,10 @@ def generate_dependent_dispatch(tup, handlers, next_call, slf, name, err, nerr):
         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]
+        codes = [
+            generate_checking_code(types[k]).sub(arg=Code(argname(k))).fill(ndb)
+            for k in relevant
+        ]
         conj = " and ".join(codes)
         if not conj:  # pragma: no cover
             # Not sure if this can happen
@@ -279,9 +260,7 @@ def generate_dependent_dispatch(tup, handlers, next_call, slf, name, err, nerr):
 
     inject["FALLTHROUGH"] = (next_call and next_call[0]) or raise_error
 
-    fn = instantiate_code(
-        symbol="__DEPENDENT_DISPATCH__", code=code, inject=inject
-    )
+    fn = instantiate_code(symbol="__DEPENDENT_DISPATCH__", code=code, inject=inject)
     return rename_function(fn, name)
 
 
@@ -311,14 +290,10 @@ class Conformer:
             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 = transfer_function(
+                func=ofn,
+                code=new_code,
             )
-            new_fn.__annotations__ = ofn.__annotations__
 
         self.ovld.register(new_fn)
 
@@ -329,90 +304,45 @@ class Conformer:
         self.code = new_code
 
 
-def rename_code(co, newname):  # pragma: no cover
-    if hasattr(co, "replace"):
-        if hasattr(co, "co_qualname"):
-            return co.replace(co_name=newname, co_qualname=newname)
-        else:
-            return co.replace(co_name=newname)
-    else:
-        return type(co)(
-            co.co_argcount,
-            co.co_kwonlyargcount,
-            co.co_nlocals,
-            co.co_stacksize,
-            co.co_flags,
-            co.co_code,
-            co.co_consts,
-            co.co_names,
-            co.co_varnames,
-            co.co_filename,
-            newname,
-            co.co_firstlineno,
-            co.co_lnotab,
-            co.co_freevars,
-            co.co_cellvars,
-        )
-
-
-def rename_function(fn, newname):
-    """Create a copy of the function with a different name."""
-    newcode = rename_code(fn.__code__, 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,
-        map_mangled,
-        code_mangled,
-    ):
+    def __init__(self, anal, special_syms, mapping):
         self.analysis = anal
-        self.recurse_sym = recurse_sym
-        self.call_next_sym = call_next_sym
-        self.ovld_mangled = ovld_mangled
-        self.map_mangled = map_mangled
-        self.code_mangled = code_mangled
+        self.syms = special_syms
+        self.mapping = mapping
+        self.ovld_mangled = mapping[recurse]
+        self.map_mangled = mapping[resolve]
+        self.code_mangled = mapping[current_code]
         self.count = count()
 
+    def is_special(self, name, *kinds):
+        return any(name in self.syms[kind] for kind in kinds)
+
     def visit_Name(self, node):
-        if node.id == self.recurse_sym:
+        if node.id in self.mapping:
             return ast.copy_location(
                 old_node=node,
-                new_node=ast.Name(self.ovld_mangled, ctx=node.ctx),
+                new_node=ast.Name(self.mapping[node.id], ctx=node.ctx),
             )
-        elif node.id == self.call_next_sym:
+        elif self.is_special(node.id, call_next):
             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,
+        if not isinstance(node.func, ast.Name) or not self.is_special(
+            node.func.id, recurse, call_next
         ):
             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
+        cn = self.is_special(node.func.id, call_next)
         tmp = f"__TMP{next(self.count)}_"
 
         def _make_lookup_call(key, arg):
             name = (
-                "__SUBTLER_TYPE"
-                if self.analysis.lookup_for(key) is subtler_type
-                else "type"
+                "__SUBTLER_TYPE" if self.analysis.lookup_for(key) is subtler_type else "type"
             )
             value = ast.NamedExpr(
                 target=ast.Name(id=f"{tmp}{key}", ctx=ast.Store()),
@@ -426,9 +356,7 @@ class NameConverter(ast.NodeTransformer):
             )
 
         # type index for positional arguments
-        type_parts = [
-            _make_lookup_call(i, arg) for i, arg in enumerate(node.args)
-        ]
+        type_parts = [_make_lookup_call(i, arg) for i, arg in enumerate(node.args)]
 
         # type index for keyword arguments
         type_parts += [
@@ -460,10 +388,7 @@ class NameConverter(ast.NodeTransformer):
         new_node = ast.Call(
             func=method,
             args=selfarg
-            + [
-                ast.Name(id=f"{tmp}{i}", ctx=ast.Load())
-                for i, arg in enumerate(node.args)
-            ],
+            + [ast.Name(id=f"{tmp}{i}", ctx=ast.Load()) for i, arg in enumerate(node.args)],
             keywords=[
                 ast.keyword(
                     arg=kw.arg,
@@ -475,37 +400,41 @@ class NameConverter(ast.NodeTransformer):
         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 _search_names(co, specials, glb, closure=None):
+    def _search(co, values, glb, closure=None):
+        if isinstance(co, CodeType):
+            if closure is not None:
+                for varname, cell in zip(co.co_freevars, closure):
+                    try:
+                        if any(cell.cell_contents is v for v in values):
+                            yield varname
+                    except ValueError:  # cell is empty
+                        pass
+            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(ct, values, glb)
+
+    return {k: list(_search(co, v, glb, closure)) for k, v in specials.items()}
 
 
 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, ovld.dispatch),
-            fn.__globals__,
-            fn.__closure__,
-        )
-    )
-    cn_syms = list(
-        _search_names(fn.__code__, (call_next,), fn.__globals__, fn.__closure__)
+    syms = _search_names(
+        fn.__code__,
+        {
+            recurse: (recurse, ovld, ovld.dispatch),
+            call_next: (call_next,),
+            resolve: (resolve,),
+            current_code: (current_code,),
+        },
+        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
-        )
+    if any(syms.values()):
+        return recode(fn, ovld, syms, newname)
     else:
         return rename_function(fn, newname)
 
@@ -536,7 +465,7 @@ def closure_wrap(tree, fname, names):
     return ast.Module(body=[wrap], type_ignores=[])
 
 
-def recode(fn, ovld, recurse_sym, call_next_sym, newname):
+def recode(fn, ovld, syms, newname):
     ovld_mangled = f"___OVLD{ovld.id}"
     map_mangled = f"___MAP{ovld.id}"
     code_mangled = f"___CODE{next(_current)}"
@@ -550,14 +479,22 @@ def recode(fn, ovld, recurse_sym, call_next_sym, newname):
             " avoid calling recurse()/call_next()"
         )
     tree = ast.parse(textwrap.dedent(src))
+
+    mapping = {
+        recurse: ovld_mangled,
+        resolve: map_mangled,
+        current_code: code_mangled,
+    }
+    for special, symbols in syms.items():
+        for sym in symbols:
+            if special in mapping:
+                mapping[sym] = mapping[special]
     new = NameConverter(
         anal=ovld.argument_analysis,
-        recurse_sym=recurse_sym,
-        call_next_sym=call_next_sym,
-        ovld_mangled=ovld_mangled,
-        map_mangled=map_mangled,
-        code_mangled=code_mangled,
+        special_syms=syms,
+        mapping=mapping,
     ).visit(tree)
+
     new.body[0].decorator_list = []
     if fn.__closure__:
         new = closure_wrap(new.body[0], "irrelevant", fn.__code__.co_freevars)
@@ -568,17 +505,14 @@ def recode(fn, ovld, recurse_sym, call_next_sym, newname):
         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
-        ]
+        [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 = transfer_function(
+        func=fn,
+        code=rename_code(new_code, newname),
+        name=newname,
+        closure=new_closure,
     )
-    new_fn.__kwdefaults__ = fn.__kwdefaults__
-    new_fn.__annotations__ = fn.__annotations__
-    new_fn = rename_function(new_fn, newname)
     new_fn.__globals__["__SUBTLER_TYPE"] = subtler_type
     new_fn.__globals__[ovld_mangled] = ovld.dispatch
     new_fn.__globals__[map_mangled] = ovld.map

ovld/signatures.py

@@ -0,0 +1,275 @@
+"""Utilities to deal with function signatures."""
+
+import inspect
+import itertools
+import typing
+from collections import OrderedDict, defaultdict
+from dataclasses import dataclass, field
+from functools import cached_property
+from types import GenericAlias
+
+from .types import normalize_type
+from .utils import MISSING, subtler_type
+
+
+class LazySignature(inspect.Signature):
+    def __init__(self, ovld):
+        super().__init__([])
+        self.ovld = ovld
+
+    def replace(
+        self, *, parameters=inspect._void, return_annotation=inspect._void
+    ):  # pragma: no cover
+        if parameters is inspect._void:
+            parameters = self.parameters.values()
+
+        if return_annotation is inspect._void:
+            return_annotation = self._return_annotation
+
+        return inspect.Signature(parameters, return_annotation=return_annotation)
+
+    @property
+    def parameters(self):
+        anal = self.ovld.analyze_arguments()
+        parameters = []
+        if anal.is_method:
+            parameters.append(
+                inspect.Parameter(
+                    name="self",
+                    kind=inspect._POSITIONAL_ONLY,
+                )
+            )
+        parameters += [
+            inspect.Parameter(
+                name=p,
+                kind=inspect._POSITIONAL_ONLY,
+            )
+            for p in anal.strict_positional_required
+        ]
+        parameters += [
+            inspect.Parameter(
+                name=p,
+                kind=inspect._POSITIONAL_ONLY,
+                default=MISSING,
+            )
+            for p in anal.strict_positional_optional
+        ]
+        parameters += [
+            inspect.Parameter(
+                name=p,
+                kind=inspect._POSITIONAL_OR_KEYWORD,
+            )
+            for p in anal.positional_required
+        ]
+        parameters += [
+            inspect.Parameter(
+                name=p,
+                kind=inspect._POSITIONAL_OR_KEYWORD,
+                default=MISSING,
+            )
+            for p in anal.positional_optional
+        ]
+        parameters += [
+            inspect.Parameter(
+                name=p,
+                kind=inspect._KEYWORD_ONLY,
+            )
+            for p in anal.keyword_required
+        ]
+        parameters += [
+            inspect.Parameter(
+                name=p,
+                kind=inspect._KEYWORD_ONLY,
+                default=MISSING,
+            )
+            for p in anal.keyword_optional
+        ]
+        return OrderedDict({p.name: p for p in parameters})
+
+
+@dataclass(frozen=True)
+class Arginfo:
+    position: typing.Optional[int]
+    name: typing.Optional[str]
+    required: bool
+    ann: type
+
+    @cached_property
+    def is_complex(self):
+        return isinstance(self.ann, GenericAlias)
+
+    @cached_property
+    def canonical(self):
+        return self.name if self.position is None else self.position
+
+
+@dataclass(frozen=True)
+class Signature:
+    types: tuple
+    return_type: type
+    req_pos: int
+    max_pos: int
+    req_names: frozenset
+    vararg: bool
+    priority: float
+    tiebreak: int = 0
+    is_method: bool = False
+    arginfo: list[Arginfo] = field(default_factory=list, hash=False, compare=False)
+
+    @classmethod
+    def extract(cls, fn):
+        typelist = []
+        sig = inspect.signature(fn)
+        max_pos = 0
+        req_pos = 0
+        req_names = set()
+        is_method = False
+
+        arginfo = []
+        for i, (name, param) in enumerate(sig.parameters.items()):
+            if name == "self" or (name == "cls" and getattr(fn, "specializer", False)):
+                if i != 0:  # pragma: no cover
+                    raise Exception(
+                        f"Argument name '{name}' marks a method and must always be in the first position."
+                    )
+                is_method = True
+                continue
+            pos = nm = None
+            ann = normalize_type(param.annotation, fn)
+            if param.kind is inspect._POSITIONAL_ONLY:
+                pos = i - is_method
+                typelist.append(ann)
+                req_pos += param.default is inspect._empty
+                max_pos += 1
+            elif param.kind is inspect._POSITIONAL_OR_KEYWORD:
+                pos = i - is_method
+                nm = param.name
+                typelist.append(ann)
+                req_pos += param.default is inspect._empty
+                max_pos += 1
+            elif param.kind is inspect._KEYWORD_ONLY:
+                nm = param.name
+                typelist.append((param.name, ann))
+                if param.default is inspect._empty:
+                    req_names.add(param.name)
+            elif param.kind is inspect._VAR_POSITIONAL:
+                raise TypeError("ovld does not support *args")
+            elif param.kind is inspect._VAR_KEYWORD:
+                raise TypeError("ovld does not support **kwargs")
+            arginfo.append(
+                Arginfo(
+                    position=pos,
+                    name=nm,
+                    required=param.default is inspect._empty,
+                    ann=ann,
+                )
+            )
+
+        return cls(
+            types=tuple(typelist),
+            return_type=normalize_type(sig.return_annotation, fn),
+            req_pos=req_pos,
+            max_pos=max_pos,
+            req_names=frozenset(req_names),
+            vararg=False,
+            is_method=is_method,
+            priority=None,
+            arginfo=arginfo,
+        )
+
+
+class ArgumentAnalyzer:
+    def __init__(self):
+        self.name_to_positions = defaultdict(set)
+        self.position_to_names = defaultdict(set)
+        self.counts = defaultdict(lambda: [0, 0])
+        self.complex_transforms = set()
+        self.total = 0
+        self.is_method = None
+        self.done = False
+
+    def add(self, fn):
+        self.done = False
+        sig = Signature.extract(fn)
+        self.complex_transforms.update(arg.canonical for arg in sig.arginfo if arg.is_complex)
+        for arg in sig.arginfo:
+            if arg.position is not None:
+                self.position_to_names[arg.position].add(arg.name)
+            if arg.name is not None:
+                self.name_to_positions[arg.name].add(arg.canonical)
+
+            cnt = self.counts[arg.canonical]
+            cnt[0] += arg.required
+            cnt[1] += 1
+
+        self.total += 1
+
+        if self.is_method is None:
+            self.is_method = sig.is_method
+        elif self.is_method != sig.is_method:  # pragma: no cover
+            raise TypeError(
+                "Some, but not all registered methods define `self`. It should be all or none."
+            )
+
+    def compile(self):
+        if self.done:
+            return
+        for name, pos in self.name_to_positions.items():
+            if len(pos) != 1:
+                if all(isinstance(p, int) for p in pos):
+                    raise TypeError(
+                        f"Argument '{name}' is declared in different positions by different methods. The same argument name should always be in the same position unless it is strictly positional."
+                    )
+                else:
+                    raise TypeError(
+                        f"Argument '{name}' is declared in a positional and keyword setting by different methods. It should be either."
+                    )
+
+        p_to_n = [list(names) for _, names in sorted(self.position_to_names.items())]
+
+        positional = list(
+            itertools.takewhile(
+                lambda names: len(names) == 1 and isinstance(names[0], str),
+                reversed(p_to_n),
+            )
+        )
+        positional.reverse()
+        strict_positional = p_to_n[: len(p_to_n) - len(positional)]
+
+        assert strict_positional + positional == p_to_n
+
+        self.strict_positional_required = [
+            f"ARG{pos + 1}"
+            for pos, _ in enumerate(strict_positional)
+            if self.counts[pos][0] == self.total
+        ]
+        self.strict_positional_optional = [
+            f"ARG{pos + 1}"
+            for pos, _ in enumerate(strict_positional)
+            if self.counts[pos][0] != self.total
+        ]
+
+        self.positional_required = [
+            names[0]
+            for pos, names in enumerate(positional)
+            if self.counts[pos + len(strict_positional)][0] == self.total
+        ]
+        self.positional_optional = [
+            names[0]
+            for pos, names in enumerate(positional)
+            if self.counts[pos + len(strict_positional)][0] != self.total
+        ]
+
+        keywords = [
+            name for _, (name,) in self.name_to_positions.items() if not isinstance(name, int)
+        ]
+        self.keyword_required = [
+            name for name in keywords if self.counts[name][0] == self.total
+        ]
+        self.keyword_optional = [
+            name for name in keywords if self.counts[name][0] != self.total
+        ]
+        self.done = True
+
+    def lookup_for(self, key):
+        return subtler_type if key in self.complex_transforms else type