<commit_before># Copyright 2016 - 2018  Ternaris.
# SPDX-License-Identifier: AGPL-3.0-only

import os
from collections.abc import Mapping, Sequence
from pickle import PickleBuffer

from capnp.lib.capnp import _DynamicEnum, _DynamicListReader, _DynamicStructReader

from marv_api.setid import SetID
from marv_api.utils import find_obj


def _to_dict(value, field=None, field_type=None, which=False):
    if isinstance(value, _DynamicStructReader):
        schema = value.schema
        dct = {}
        for name in schema.non_union_fields:
            dct[name] = _to_dict(getattr(value, name), field=schema.fields[name], which=which)

        if schema.union_fields:
            _which = value.which()
            dct[_which] = _to_dict(getattr(value, _which), field=schema.fields[_which], which=which)
            if which:
                dct['_which'] = _which

        return dct

    if isinstance(value, _DynamicListReader):
        element_type = (field_type or field.proto.slot.type).list.elementType
        return [
            _to_dict(element, field_type=element_type, which=which) for element in value
        ]

    if isinstance(value, _DynamicEnum):
        return value._as_str()  # pylint: disable=protected-access

    return value


def _wrap(value, streamdir, setdir, field=None, field_type=None):
    if isinstance(value, _DynamicStructReader):
        return Wrapper(value, streamdir, setdir)

    if isinstance(value, _DynamicListReader):
        if field:
            element_type = field.proto.slot.type.list.elementType
        else:
            element_type = field_type.list.elementType
        return ListWrapper(value, element_type, streamdir, setdir)

    return value


class ListWrapper:
    def __init__(self, list_reader, field_type, streamdir, setdir):
        assert isinstance(list_reader, _DynamicListReader), type(list_reader)
        self._field_type = field_type
        self._reader = list_reader
        self._streamdir = streamdir
        self._setdir = setdir

    def _wrap(self, item):
        return _wrap(item, self._streamdir, self._setdir, field_type=self._field_type)

    def __getitem__(self, idx):
        if isinstance(idx, slice):  # pycapnp can't handle slices
            return [self._wrap(x) for x in list(self._reader)[idx]]
        return self._wrap(self._reader[idx])

    def __eq__(self, other):
        return self[:] == other

    def __iter__(self):
        return (self._wrap(x) for x in self._reader)

    def __len__(self):
        return len(self._reader)

    def __repr__(self):
        return f"[{', '.join([repr(x) for x in self])}]"


class Wrapper:
    def __init__(self, struct_reader, streamdir, setdir, storedir=None):
        assert isinstance(struct_reader, _DynamicStructReader), type(struct_reader)
        self._reader = struct_reader
        self._streamdir = streamdir  # HACK: overloaded
        self._setdir = setdir
        self._storedir = storedir  # HACK: patched by compare

    # @property
    # def path(self):  # HACK: overload
    #     return os.path.realpath(os.path.join(self._streamdir, self.name))

    def __reduce_ex__(self, protocol):
        if protocol < 5:
            raise RuntimeError('Needs at least pickle protocol 5')

        builder = self._reader.as_builder()
        protoname = builder.schema.node.displayName.replace('/', '.')\
                                                   .replace('.capnp:', '_capnp:')
        meta = {
            'protoname': protoname,
            'streamdir': self._streamdir,
            'setdir': self._setdir,
            'storedir': self._storedir,
        }
        segments = builder.to_segments()
        return (self.from_segments, (meta, *[PickleBuffer(x) for x in segments]))

    @property
    def relpath(self):  # HACK: overload
        path = os.path.relpath(self.path, self._setdir)
        return path.lstrip('.')

    @classmethod
    def from_segments(cls, meta, *segments):
        schema = find_obj(meta.pop('protoname'))
        struct_reader = schema.from_segments(segments)
        return cls(struct_reader, **meta)

    @classmethod
    def from_dict(cls, schema, data):
        setid = data.pop('id', None)
        if isinstance(setid, SetID):
            data['id0'], data['id1'] = setid.lohi
        dct = cls._unwrap(data)
        struct_reader = schema.new_message(**dct).as_reader()
        return cls(struct_reader, None, None)

    def to_dict(self, which=None):
        return _to_dict(self._reader, which=which)

    def load(self, node):
        from marv_store import Store  # pylint: disable=import-outside-toplevel
        store = Store(self._storedir, {node.name: node})
        return store.load(self._setdir, node, default=None)

    @classmethod
    def _unwrap(cls, data):
        unwrap = cls._unwrap
        if isinstance(data, Mapping):
            return {k: unwrap(v) for k, v in data.items()}

        if isinstance(data, Sequence) and not isinstance(data, (bytes, str)):
            return [unwrap(x) for x in data]

        if isinstance(data, Wrapper):
            return data._reader  # pylint: disable=protected-access

        return data

    def __getattr__(self, name):
        parts = name.split('_')
        field_name = parts[0] + ''.join(((x[0].upper() + x[1:]) if x else '_') for x in parts[1:])

        # pylint: disable=protected-access
        if field_name in self._reader.schema.fieldnames and self._reader._has(field_name):
            field = self._reader.schema.fields[field_name]
            value = getattr(self._reader, name)

            if self._streamdir and name == 'path' and \
               self._reader.schema.node.displayName == 'marv_nodes/types.capnp:File' and \
               not os.path.exists(value):
                relpath = os.path.relpath(value, self._setdir)
                if relpath[0] == '.':
                    return os.path.join(self._setdir, relpath[1:])
<commit_msg>Enable getattr hooks for pycapnp wrapper<commit_after># Copyright 2016 - 2018  Ternaris.
# SPDX-License-Identifier: AGPL-3.0-only

import os
from collections.abc import Mapping, Sequence
from pickle import PickleBuffer

from capnp.lib.capnp import _DynamicEnum, _DynamicListReader, _DynamicStructReader

from marv_api.setid import SetID
from marv_api.utils import find_obj

GETATTR_HOOKS = []


def _to_dict(value, field=None, field_type=None, which=False):
    if isinstance(value, _DynamicStructReader):
        schema = value.schema
        dct = {}
        for name in schema.non_union_fields:
            dct[name] = _to_dict(getattr(value, name), field=schema.fields[name], which=which)

        if schema.union_fields:
            _which = value.which()
            dct[_which] = _to_dict(getattr(value, _which), field=schema.fields[_which], which=which)
            if which:
                dct['_which'] = _which

        return dct

    if isinstance(value, _DynamicListReader):
        element_type = (field_type or field.proto.slot.type).list.elementType
        return [
            _to_dict(element, field_type=element_type, which=which) for element in value
        ]

    if isinstance(value, _DynamicEnum):
        return value._as_str()  # pylint: disable=protected-access

    return value


def _wrap(value, streamdir, setdir, field=None, field_type=None):
    if isinstance(value, _DynamicStructReader):
        return Wrapper(value, streamdir, setdir)

    if isinstance(value, _DynamicListReader):
        if field:
            element_type = field.proto.slot.type.list.elementType
        else:
            element_type = field_type.list.elementType
        return ListWrapper(value, element_type, streamdir, setdir)

    return value


class ListWrapper:
    def __init__(self, list_reader, field_type, streamdir, setdir):
        assert isinstance(list_reader, _DynamicListReader), type(list_reader)
        self._field_type = field_type
        self._reader = list_reader
        self._streamdir = streamdir
        self._setdir = setdir

    def _wrap(self, item):
        return _wrap(item, self._streamdir, self._setdir, field_type=self._field_type)

    def __getitem__(self, idx):
        if isinstance(idx, slice):  # pycapnp can't handle slices
            return [self._wrap(x) for x in list(self._reader)[idx]]
        return self._wrap(self._reader[idx])

    def __eq__(self, other):
        return self[:] == other

    def __iter__(self):
        return (self._wrap(x) for x in self._reader)

    def __len__(self):
        return len(self._reader)

    def __repr__(self):
        return f"[{', '.join([repr(x) for x in self])}]"


class Wrapper:
    def __init__(self, struct_reader, streamdir, setdir, storedir=None):
        assert isinstance(struct_reader, _DynamicStructReader), type(struct_reader)
        self._reader = struct_reader
        self._streamdir = streamdir  # HACK: overloaded
        self._setdir = setdir
        self._storedir = storedir  # HACK: patched by compare

    # @property
    # def path(self):  # HACK: overload
    #     return os.path.realpath(os.path.join(self._streamdir, self.name))

    def __reduce_ex__(self, protocol):
        if protocol < 5:
            raise RuntimeError('Needs at least pickle protocol 5')

        builder = self._reader.as_builder()
        protoname = builder.schema.node.displayName.replace('/', '.')\
                                                   .replace('.capnp:', '_capnp:')
        meta = {
            'protoname': protoname,
            'streamdir': self._streamdir,
            'setdir': self._setdir,
            'storedir': self._storedir,
        }
        segments = builder.to_segments()
        return (self.from_segments, (meta, *[PickleBuffer(x) for x in segments]))

    @property
    def relpath(self):  # HACK: overload
        path = os.path.relpath(self.path, self._setdir)
        return path.lstrip('.')

    @classmethod
    def from_segments(cls, meta, *segments):
        schema = find_obj(meta.pop('protoname'))
        struct_reader = schema.from_segments(segments)
        return cls(struct_reader, **meta)

    @classmethod
    def from_dict(cls, schema, data):
        setid = data.pop('id', None)
        if isinstance(setid, SetID):
            data['id0'], data['id1'] = setid.lohi
        dct = cls._unwrap(data)
        struct_reader = schema.new_message(**dct).as_reader()
        return cls(struct_reader, None, None)

    def to_dict(self, which=None):
        return _to_dict(self._reader, which=which)

    def load(self, node):
        from marv_store import Store  # pylint: disable=import-outside-toplevel
        store = Store(self._storedir, {node.name: node})
        return store.load(self._setdir, node, default=None)

    @classmethod
    def _unwrap(cls, data):
        unwrap = cls._unwrap
        if isinstance(data, Mapping):
            return {k: unwrap(v) for k, v in data.items()}

        if isinstance(data, Sequence) and not isinstance(data, (bytes, str)):
            return [unwrap(x) for x in data]

        if isinstance(data, Wrapper):
            return data._reader  # pylint: disable=protected-access

        return data

    def __getattr__(self, name):  # noqa: C901  # pylint: disable=too-many-return-statements
        for hook in GETATTR_HOOKS:
            handled, ret = hook(self, self._reader.schema.node.displayName, name)
            if handled:
                return ret

        parts = name.split('_')
        field_name = parts[0] + ''.join(((x[0].upper() + x[1:]) if x else '_') for x in parts[1:])

        # pylint: disable=protected-access
        if field_name in self._reader.schema.fieldnames and self._reader._has(field_name):
            field = self._reader.schema.fields[field_name]
            value = getattr(self._reader, name)

            if self._streamdir and name == 'path' and \
               self._reader.schema.node.displayName == 'marv_nodes/types.capnp:File' and \
               not os.path.exists(value):
                relpath = os.path.relpath(value, self._setdir)
                if relpath[0] == '.':
                    return os.path.join(self._setdir, relpath[1:])

{'fraction_non_alphanumeric': 0.10994721214120752, 'fraction_numerical': 0.0035466842626195975, 'mean_word_length': 2.5661764705882355, 'pattern_counts': {'":': 0, '<': 5, '<?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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1198248', 'n_tokens_mistral': 3726, 'n_tokens_neox': 3622, 'n_words': 865}