<commit_before>from contextlib import contextmanager
from collections import deque, OrderedDict

from .. import graph
from ..types import Sequence, SequenceMeta, Record, RecordMeta, Optional
from ..types import MappingMeta, OptionalMeta, Unknown, UnknownMeta
from ..types import TypeRef, TypeRefMeta

from .refs import NamedRef, Ref
from .nodes import NodeTransformer, Symbol, Keyword, Tuple, List


class GraphTypes(graph.GraphVisitor):

    def visit_graph(self, obj):
        types = OrderedDict((node.name, self.visit(node)) for node in obj.nodes)
        types.update(self.visit(obj.root).__field_types__)
        return types

    def visit_node(self, obj):
        return Record[[(f.name, self.visit(f)) for f in obj.fields]]

    def visit_link(self, obj):
        if obj.type_enum is graph.Maybe:
            return Optional[TypeRef[obj.node]]
        elif obj.type_enum is graph.One:
            return TypeRef[obj.node]
        elif obj.type_enum is graph.Many:
            return Sequence[TypeRef[obj.node]]
        else:
            raise TypeError(repr(obj.type_enum))

    def visit_field(self, obj):
        return obj.type or Unknown


def graph_types(graph_):
    return GraphTypes().visit(graph_)


def fn_types(functions):
    return {fn.__def_name__: fn.__def_type__ for fn in functions}


class Environ(object):

    def __init__(self, values):
        env_vars = {}
        for name, value in values.items():
            env_vars[name] = NamedRef(None, name, value)
        self.vars = deque([env_vars])

    @contextmanager
    def push(self, mapping):
        self.vars.append(mapping)
        try:
            yield
        finally:
            self.vars.pop()

    def __getitem__(self, key):
        for d in reversed(self.vars):
            try:
                return d[key]
            except KeyError:
                continue
        else:
            raise KeyError(repr(key))

    def __contains__(self, key):
        return any(key in d for d in self.vars)


def get_type(types, obj):
    if isinstance(obj, TypeRefMeta):
        return types[obj.__type_name__]
    else:
        return obj


def node_type(types, node):
    return get_type(types, node.__ref__.to)


def check_type(types, t1, t2):
    t1 = get_type(types, t1)
    t2 = get_type(types, t2)
    if isinstance(t2, UnknownMeta):
        pass
    else:
        if isinstance(t1, type(t2)):
            if isinstance(t2, OptionalMeta):
                check_type(types, t1.__type__, t2.__type__)
            elif isinstance(t2, SequenceMeta):
                check_type(types, t1.__item_type__, t2.__item_type__)
            elif isinstance(t2, MappingMeta):
                check_type(types, t1.__key_type__, t2.__key_type__)
                check_type(types, t1.__value_type__, t2.__value_type__)
            elif isinstance(t2, RecordMeta):
                for key, v2 in t2.__field_types__.items():
                    v1 = t1.__field_types__.get(key)
                    if v1 is None:
                        raise TypeError('Missing field {}'.format(key))
                    v1 = get_type(types, v1)
                    check_type(types, v1, v2)
        else:
            raise TypeError('Types mismatch: {} != {}'
                            .format(type(t1), type(t2)))


class Checker(NodeTransformer):

    def __init__(self, types):
        self.types = types
        self.env = Environ(types)

    def visit_get_expr(self, node):
        sym, obj, name = node.values
        assert isinstance(name, Symbol), type(name)
        obj = self.visit(obj)
        assert hasattr(obj, '__ref__'), 'Object does not have a reference'

        ref_to = node_type(self.types, obj)<commit_msg>Improved TypeError's message when types are not compatible<commit_after>from contextlib import contextmanager
from collections import deque, OrderedDict

from .. import graph
from ..types import Sequence, SequenceMeta, Record, RecordMeta, Optional
from ..types import MappingMeta, OptionalMeta, Unknown, UnknownMeta
from ..types import TypeRef, TypeRefMeta

from .refs import NamedRef, Ref
from .nodes import NodeTransformer, Symbol, Keyword, Tuple, List


class GraphTypes(graph.GraphVisitor):

    def visit_graph(self, obj):
        types = OrderedDict((node.name, self.visit(node)) for node in obj.nodes)
        types.update(self.visit(obj.root).__field_types__)
        return types

    def visit_node(self, obj):
        return Record[[(f.name, self.visit(f)) for f in obj.fields]]

    def visit_link(self, obj):
        if obj.type_enum is graph.Maybe:
            return Optional[TypeRef[obj.node]]
        elif obj.type_enum is graph.One:
            return TypeRef[obj.node]
        elif obj.type_enum is graph.Many:
            return Sequence[TypeRef[obj.node]]
        else:
            raise TypeError(repr(obj.type_enum))

    def visit_field(self, obj):
        return obj.type or Unknown


def graph_types(graph_):
    return GraphTypes().visit(graph_)


def fn_types(functions):
    return {fn.__def_name__: fn.__def_type__ for fn in functions}


class Environ(object):

    def __init__(self, values):
        env_vars = {}
        for name, value in values.items():
            env_vars[name] = NamedRef(None, name, value)
        self.vars = deque([env_vars])

    @contextmanager
    def push(self, mapping):
        self.vars.append(mapping)
        try:
            yield
        finally:
            self.vars.pop()

    def __getitem__(self, key):
        for d in reversed(self.vars):
            try:
                return d[key]
            except KeyError:
                continue
        else:
            raise KeyError(repr(key))

    def __contains__(self, key):
        return any(key in d for d in self.vars)


def get_type(types, obj):
    if isinstance(obj, TypeRefMeta):
        return types[obj.__type_name__]
    else:
        return obj


def node_type(types, node):
    return get_type(types, node.__ref__.to)


def check_type(types, t1, t2):
    t1 = get_type(types, t1)
    t2 = get_type(types, t2)
    if isinstance(t2, UnknownMeta):
        pass
    else:
        if isinstance(t1, type(t2)):
            if isinstance(t2, OptionalMeta):
                check_type(types, t1.__type__, t2.__type__)
            elif isinstance(t2, SequenceMeta):
                check_type(types, t1.__item_type__, t2.__item_type__)
            elif isinstance(t2, MappingMeta):
                check_type(types, t1.__key_type__, t2.__key_type__)
                check_type(types, t1.__value_type__, t2.__value_type__)
            elif isinstance(t2, RecordMeta):
                for key, v2 in t2.__field_types__.items():
                    v1 = t1.__field_types__.get(key)
                    if v1 is None:
                        raise TypeError('Missing field {}'.format(key))
                    v1 = get_type(types, v1)
                    check_type(types, v1, v2)
        else:
            raise TypeError('Types mismatch, {!r} != {!r}'.format(t1, t2))


class Checker(NodeTransformer):

    def __init__(self, types):
        self.types = types
        self.env = Environ(types)

    def visit_get_expr(self, node):
        sym, obj, name = node.values
        assert isinstance(name, Symbol), type(name)
        obj = self.visit(obj)
        assert hasattr(obj, '__ref__'), 'Object does not have a reference'

        ref_to = node_type(self.types, obj)
{'fraction_non_alphanumeric': 0.09945872801082543, 'fraction_numerical': 0.008660351826792964, 'mean_word_length': 2.266018559434379, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4851833', 'n_tokens_mistral': 2257, 'n_tokens_neox': 2171, 'n_words': 581}