fny
fny

Reputation: 33587

Using Decorators to Track and Return Methods

I'm trying to define a decorator that will mark certain methods, and then later provide the ability to call all marked methods by method.

This following should given an idea how I'd like everything to behave, but isn't functional for obvious reasons

def mark(method):
    # mark_methods should belong to the class wherein methods are being marked
    mark_methods.append(method)
    return method

class BaseClass():
    # This should belong to each subclasses independently
    marked_methods = []

    def bundle(self):
        data = {}
        for marked_method in Person.marked_methods:
            data[marked_method.__name___] = marked_method()
        return data

class Person(BaseClass):
    @mark
    def name(self):
        return 'Fred'

    def say_hi(self):
        return 'Hi'

class Dog(BaseClass):
    @mark
    def name(self):
        return 'Fido'

    @mark
    def fur_color(self):
        return 'Black'

    def bark(self):
        return 'Woof'

Person.marked_methods # => [name]
Person().bundle() # => {'name': 'Fred'}

Dog.marked_methods # => [name, fur_color]
Dog().bundle() # => {'name': 'Fido', 'fur_color': 'Black'}

Ideally, this behavior would be wrapped in a class that could be inherited by other classes too.

Here's a version of what I'm looking for that exhibits a similar behavior without using decorators. Instead it rely on manipulating functions that start with a certain prefix mark_whatever:

MARK_PREFIX = 'mark_'

class TrackingWithoutDecoratorClass():
    @classmethod
    def __init_subclass__(cls, **kwargs):
        """Tracks all methods added on init
        """

        # Mapping of transformer names to transformer functions
        cls.marked = {}

        for name, func in vars(cls).items():
            # Add all functions starting with the `MARK_PREFIX` to the
            # marked registry
            if name.startswith(MARK_PREFIX):
                registry_name = name.replace(MARK_PREFIX, '')
                cls.marked[registry_name] = func


class Example(TrackingWithoutDecoratorClass):
  def mark_one(self):
    return 1
  def mark_two(self):
    return 2
  def not_marked_three(self):
    return 3

print(Example.marked.keys())

Here's a working REPL.

Upvotes: 0

Views: 67

Answers (2)

fny
fny

Reputation: 33587

The trick is to use a metaclass to store the decorated methods once the target classes inherit from a base class:

def mark(method):
    method.marked = True
    return method

class MarkTracking(type):
    def __new__(cls, name, bases, attr):
        marked = []
        for obj in attr.values():
            if hasattr(obj, 'marked'):
                marked.append(obj)
        attr['marked_methods'] = marked
        return type.__new__(cls, name, bases, attr)

class BaseClass(metaclass=MarkTracking):
    def bundle(self):
        data = {}
        for marked_method in self.__class__.marked_methods:
            data[marked_method.__name__] = marked_method(self)
        return data

class Person(BaseClass):
    @mark
    def name(self):
        return 'Fred'

    def say_hi(self):
        return 'Hi'

class Dog(BaseClass):
    @mark
    def name(self):
        return 'Fido'

    @mark
    def fur_color(self):
        return 'Black'

    def bark(self):
        return 'Woof'

print(Person.marked_methods) # => [name]
print(Person().bundle()) # => {'name': 'Fred'}

print(Dog.marked_methods) # => [name, fur_color]
print(Dog().bundle()) # => {'name': 'Fido', 'fur_color': 'Black'}

Upvotes: 0

Danil Speransky
Danil Speransky

Reputation: 30463

I would suggest to mark the methods themselves:

import inspect

def mark(func):
    func.marked = True
    return func

class Base():
    @classmethod
    def marked_methods(cls):
        return [n for n, f in inspect.getmembers(cls) if hasattr(f, 'marked')]

    def bundle(self):
        return {m: getattr(self, m)() for m in self.marked_methods()}

class Person(Base):
    @mark
    def name(self):
        return 'Fred'

    def say(self):
        return 'Hi'

class Dog(Base):
    @mark
    def name(self):
        return 'Fido'

    def bark(self):
        return 'Woof'

    @mark
    def color(self):
        return 'Black'

print(Person.marked_methods()) #=> ['name']
print(Person().bundle())       #=> {'name': 'Fred'}
print(Dog.marked_methods())    #=> ['color', 'name']
print(Dog().bundle())          #=> {'color': 'Black', 'name': 'Fido'}

Upvotes: 1

Related Questions