Source code for rootpy.tree.treemodel

from __future__ import absolute_import, print_function

import sys
import inspect
import types
if sys.version_info[0] < 3:
    from cStringIO import StringIO
else:
    from io import StringIO

import ROOT

from .. import log; log = log[__name__]
from .treetypes import Column
from .treebuffer import TreeBuffer

__all__ = [
    'TreeModel',
]


class TreeModelMeta(type):
    """
    Metaclass for all TreeModels
    Addition/subtraction of TreeModels is handled
    as set union and difference of class attributes
    """
    def __new__(cls, name, bases, dct):
        for attr, value in dct.items():
            TreeModelMeta.checkattr(attr, value)
        return type.__new__(cls, name, bases, dct)

    def __add__(cls, other):
        return type('_'.join([cls.__name__, other.__name__]),
                    (cls, other), {})

    def __iadd__(cls, other):
        return cls.__add__(other)

    def __sub__(cls, other):
        attrs = dict(set(cls.get_attrs()).difference(set(other.get_attrs())))
        return type('_'.join([cls.__name__, other.__name__]),
                    (TreeModel,), attrs)

    def __isub__(cls, other):
        return cls.__sub__(other)

    def __setattr__(cls, attr, value):
        TreeModelMeta.checkattr(attr, value)
        type.__setattr__(cls, attr, value)

    @classmethod
    def checkattr(metacls, attr, value):
        """
        Only allow class attributes that are instances of
        rootpy.types.Column, ROOT.TObject, or ROOT.ObjectProxy
        """
        if not isinstance(value, (
                types.MethodType,
                types.FunctionType,
                classmethod,
                staticmethod,
                property)):
            if attr in dir(type('dummy', (object,), {})) + \
                    ['__metaclass__', '__qualname__']:
                return
            if attr.startswith('_'):
                raise SyntaxError(
                    "TreeModel attribute `{0}` "
                    "must not start with `_`".format(attr))
            if not inspect.isclass(value):
                if not isinstance(value, Column):
                    raise TypeError(
                        "TreeModel attribute `{0}` "
                        "must be an instance of "
                        "`rootpy.tree.treetypes.Column`".format(attr))
                return
            if not issubclass(value, (ROOT.TObject, ROOT.ObjectProxy)):
                raise TypeError(
                    "TreeModel attribute `{0}` must inherit "
                    "from `ROOT.TObject` or `ROOT.ObjectProxy`".format(
                        attr))

    def prefix(cls, name):
        """
        Create a new TreeModel where class attribute
        names are prefixed with ``name``
        """
        attrs = dict([(name + attr, value) for attr, value in cls.get_attrs()])
        return TreeModelMeta(
            '_'.join([name, cls.__name__]),
            (TreeModel,), attrs)

    def suffix(cls, name):
        """
        Create a new TreeModel where class attribute
        names are suffixed with ``name``
        """
        attrs = dict([(attr + name, value) for attr, value in cls.get_attrs()])
        return TreeModelMeta(
            '_'.join([cls.__name__, name]),
            (TreeModel,), attrs)

    def get_attrs(cls):
        """
        Get all class attributes ordered by definition
        """
        ignore = dir(type('dummy', (object,), {})) + ['__metaclass__']
        attrs = [
            item for item in inspect.getmembers(cls) if item[0] not in ignore
            and not isinstance(
                item[1], (
                    types.FunctionType,
                    types.MethodType,
                    classmethod,
                    staticmethod,
                    property))]
        # sort by idx and use attribute name to break ties
        attrs.sort(key=lambda attr: (getattr(attr[1], 'idx', -1), attr[0]))
        return attrs

    def to_struct(cls, name=None):
        """
        Convert the TreeModel into a compiled C struct
        """
        if name is None:
            name = cls.__name__
        basic_attrs = dict([(attr_name, value)
                            for attr_name, value in cls.get_attrs()
                            if isinstance(value, Column)])
        if not basic_attrs:
            return None
        src = 'struct {0} {{'.format(name)
        for attr_name, value in basic_attrs.items():
            src += '{0} {1};'.format(value.type.typename, attr_name)
        src += '};'
        if ROOT.gROOT.ProcessLine(src) != 0:
            return None
        return getattr(ROOT, name, None)

    def __repr__(cls):
        out = StringIO()
        for name, value in cls.get_attrs():
            print('{0} -> {1}'.format(name, value), file=out)
        return out.getvalue()[:-1]

    def __str__(cls):
        return repr(cls)


# TreeModel.__new__
def __new__(cls):
    """
    Return a TreeBuffer for this TreeModel
    """
    treebuffer = TreeBuffer()
    for name, attr in cls.get_attrs():
        treebuffer[name] = attr()
    return treebuffer


# metaclass syntax compatible with both Python 2 and Python 3
TreeModel = TreeModelMeta('TreeModel', (object,), {'__new__': __new__})