Skip to content
This repository
tree: ae56ec7040
Fetching contributors…

Octocat-spinner-32-eaf2f5

Cannot retrieve contributors at this time

file 81 lines (56 sloc) 1.966 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
import warnings


TYPES = {}


class register(object):

    def __init__(self, names=None, demote=None, builtin=False):

        self.names = names
        if names is not None:
            if type(names) not in (list, tuple):
                raise TypeError("names must be a list or tuple")
        self.demote = demote
        self.builtin = builtin

    def __call__(self, cls):

        init_methods = []

        if not self.builtin:
            if hasattr(cls, '_post_init'):
                init_methods.append(cls._post_init)

            # all rootpy classes which inherit from ROOT classes
            # must place the ROOT base class as the last class
            # in the inheritance list
            rootbase = cls.__bases__[-1]
            cls_names = [rootbase.__name__]
        else:
            cls_names = [cls.__name__]

        if self.names is not None:
            cls_names += self.names

        cls_names_up = list(set([name.upper() for name in cls_names]))

        for name in cls_names_up:
            if name in TYPES:
                warnings.warn("Duplicate registration of class %s" % name)

            TYPES[name] = {
                'class': cls,
                'init': init_methods,
                'demote': self.demote
            }

        return cls


def lookup(cls):

    rootpy_cls, inits = lookup_by_name(cls.__name__)
    if rootpy_cls is None:
        # ROOT class not registered, pass-through
        return cls, []
    return rootpy_cls, inits


def lookup_by_name(cls_name):

    cls_name = cls_name.upper()
    if cls_name in TYPES:
        entry = TYPES[cls_name]
        return entry['class'], entry['init']
    # ROOT class not registered...
    return None, []


def lookup_demotion(cls):

    cls_name = cls.__name__.upper()
    if cls_name in TYPES:
        entry = TYPES[cls_name]
        demote = entry['demote']
        if demote is None:
            return cls_name
        return demote
    return None
Something went wrong with that request. Please try again.