pypy / pypy / annotation /

from __future__ import absolute_import

import types
from pypy.annotation.model import SomeBool, SomeInteger, SomeString,\
     SomeFloat, SomeList, SomeDict, s_None, \
     SomeObject, SomeInstance, SomeTuple, lltype_to_annotation,\
     unionof, SomeUnicodeString, SomeType
from pypy.annotation.listdef import ListDef
from pypy.annotation.dictdef import DictDef

_annotation_cache = {}

def _annotation_key(t):
    from pypy.rpython import extregistry
    if type(t) is list:
        assert len(t) == 1
        return ('list', _annotation_key(t[0]))
    elif type(t) is dict:
        assert len(t.keys()) == 1
        return ('dict', _annotation_key(t.items()[0]))
    elif isinstance(t, tuple):
        return tuple([_annotation_key(i) for i in t])
    elif extregistry.is_registered(t):
        # XXX should it really be always different?
        return t
    return t

def annotation(t, bookkeeper=None):
    if bookkeeper is None:
        key = _annotation_key(t)
            return _annotation_cache[key]
        except KeyError:
            t = _compute_annotation(t, bookkeeper)
            _annotation_cache[key] = t
            return t
    return _compute_annotation(t, bookkeeper)

def _compute_annotation(t, bookkeeper=None):
    from pypy.rpython.lltypesystem import lltype
    from pypy.rpython import extregistry
    if isinstance(t, SomeObject):
        return t
    elif isinstance(t, lltype.LowLevelType):
        return lltype_to_annotation(t)
    elif isinstance(t, list):
        assert len(t) == 1, "We do not support type joining in list"
        listdef = ListDef(bookkeeper, annotation(t[0]), mutated=True, resized=True)
        return SomeList(listdef)
    elif isinstance(t, tuple):
        return SomeTuple(tuple([annotation(i) for i in t]))
    elif isinstance(t, dict):
        assert len(t) == 1, "We do not support type joining in dict"
        result = SomeDict(DictDef(bookkeeper, annotation(t.keys()[0]),
        return result
    elif type(t) is types.NoneType:
        return s_None
    elif extregistry.is_registered(t):
        entry = extregistry.lookup(t)
        entry.bookkeeper = bookkeeper
        return entry.compute_result_annotation()
        return annotationoftype(t, bookkeeper)

def annotationoftype(t, bookkeeper=False):
    from pypy.rpython import extregistry

    """The most precise SomeValue instance that contains all
    objects of type t."""
    assert isinstance(t, (type, types.ClassType))
    if t is bool:
        return SomeBool()
    elif t is int:
        return SomeInteger()
    elif t is float:
        return SomeFloat()
    elif issubclass(t, str): # py.lib uses annotated str subclasses
        return SomeString()
    elif t is unicode:
        return SomeUnicodeString()
    elif t is types.NoneType:
        return s_None
    elif bookkeeper and extregistry.is_registered_type(t, bookkeeper.policy):
        entry = extregistry.lookup_type(t, bookkeeper.policy)
        return entry.compute_annotation_bk(bookkeeper)
    elif t is type:
        return SomeType()
    elif bookkeeper and not hasattr(t, '_freeze_'):
        classdef = bookkeeper.getuniqueclassdef(t)
        return SomeInstance(classdef)
        raise AssertionError("annotationoftype(%r)" % (t,))

class Sig(object):

    def __init__(self, *argtypes):
        self.argtypes = argtypes
    def __call__(self, funcdesc, inputcells):
        from pypy.rpython.lltypesystem import lltype
        args_s = []
        from pypy.annotation import model as annmodel
        for i, argtype in enumerate(self.argtypes):
            if isinstance(argtype, (types.FunctionType, types.MethodType)):
                argtype = argtype(*inputcells)
            if isinstance(argtype, lltype.LowLevelType) and\
                argtype is lltype.Void:
                # XXX the mapping between Void and annotation
                # is not quite well defined
                s_input = inputcells[i]
                assert isinstance(s_input, annmodel.SomePBC)
                assert s_input.is_constant()
            elif argtype is None:
                args_s.append(inputcells[i])     # no change
                args_s.append(annotation(argtype, bookkeeper=funcdesc.bookkeeper))
        if len(inputcells) != len(args_s):
            raise Exception("%r: expected %d args, got %d" % (funcdesc,
        for i, (s_arg, s_input) in enumerate(zip(args_s, inputcells)):
            s_input = unionof(s_input, s_arg)
            if not s_arg.contains(s_input):
                raise Exception("%r argument %d:\n"
                                "expected %s,\n"
                                "     got %s" % (funcdesc, i+1,
        inputcells[:] = args_s

def finish_type(paramtype, bookkeeper, func):
    from pypy.annotation.types import SelfTypeMarker
    if isinstance(paramtype, SomeObject):
        return paramtype
    elif isinstance(paramtype, SelfTypeMarker):
        raise Exception("%r argument declared as annotation.types.self(); class needs decorator rlib.signature.finishsigs()" % (func,))
        return paramtype(bookkeeper)

def enforce_signature_args(funcdesc, paramtypes, actualtypes):
    assert len(paramtypes) == len(actualtypes)
    params_s = [finish_type(paramtype, funcdesc.bookkeeper, funcdesc.pyobj) for paramtype in paramtypes]
    for i, (s_param, s_actual) in enumerate(zip(params_s, actualtypes)):
        if not s_param.contains(s_actual):
            raise Exception("%r argument %d:\n"
                            "expected %s,\n"
                            "     got %s" % (funcdesc, i+1, s_param, s_actual))
    actualtypes[:] = params_s

def enforce_signature_return(funcdesc, sigtype, inferredtype):
    return finish_type(sigtype, funcdesc.bookkeeper, funcdesc.pyobj)