Kenji Noguchi
Kenji Noguchi

Reputation: 1773

how to collect class variables while preserving the order

I'm trying to create schema classes similar to SQLAlchemy declarative_base. The schema classes should be extensible using inheritance. Here is my attempt so far. It works fine except the properties (class variables) are sorted in the alphabetical order.

import inspect
class BaseProp(object):
    def __init__(self, name):
        self.name = name
class StringProp(BaseProp): pass
class EmailProp(BaseProp): pass

class BaseSchema(object):
    @classmethod
    def _collect_props(cls):
        prop_filter =  lambda a:not(inspect.isroutine(a)) and issubclass(a.__class__, BaseProp)
        return inspect.getmembers(cls, prop_filter)

    @classmethod
    def as_list(cls):
        return cls._collect_props()

    @classmethod
    def as_dict(cls):
        return dict(cls._collect_props())

class UserIdentSchema(BaseSchema):
    email = EmailProp('email')
    password = StringProp('password')

class UserProfileSchema(UserIdentSchema):
    name = StringProp('name')
    phone = StringProp('phone')

from pprint import pprint
pprint(UserProfileSchema.as_list())

Here is the result. Notice the properties are sorted by the alphabetical order.

[('email', <__main__.EmailProp object at 0x10518a950>),
 ('name', <__main__.StringProp object at 0x10517d910>),
 ('password', <__main__.StringProp object at 0x10517d8d0>),
 ('phone', <__main__.StringProp object at 0x10517d950>)]

I want to have the base schema props at the top, then props from the subclasses. What's the best way to achieve this? Do I have to go through AST...?

EDIT: I need to keep the order of properties within a class as well.

Upvotes: 1

Views: 136

Answers (1)

mgilson
mgilson

Reputation: 309899

So here's a working solution -- I don't guarantee that this is production worthy code ... e.g. I haven't thought too much about what would happen if you tried to mix in another (non-BaseSchema) class with a BaseSchema subclass. I think that might work, but you'd need to try it out to see ...

import ast
import inspect
from collections import OrderedDict

class _NodeTagger(ast.NodeVisitor):
    def __init__(self):
        self.class_attribute_names = {}

    def visit_Assign(self, node):
        for target in node.targets:
            self.class_attribute_names[target.id] = target.lineno

    # Don't visit Assign nodes inside Function Definitions.
    def visit_FunctionDef(self, unused_node):
        return None


class BaseProp(object):
    def __init__(self, name):
        self.name = name


class StringProp(BaseProp): pass


class EmailProp(BaseProp): pass


class _BaseSchemaType(type):
    def __init__(cls, name, bases, dct):
        cls._properties = OrderedDict()
        for b in bases:
          # if the object has a metaclass which is this
          # class (or subclass of this class...)
          # then we add it's properties to our own.
          if issubclass(type(b), _BaseSchemaType):
            cls._properties.update(b._properties)

        # Finally, we add our own properties.  We find our own source code
        # read it and tag the nodes where we find assignment.
        source = inspect.getsource(cls)
        nt = _NodeTagger()
        nt.visit(ast.parse(source))
        attr_names = nt.class_attribute_names
        properties = {n: prop for n, prop in dct.items()
                      if isinstance(prop, BaseProp)}
        sorted_attrs = sorted(properties.items(),
                             key=lambda item: attr_names[item[0]])
        cls._properties.update(sorted_attrs)

    # methods on a metaclass are basically just classmethods ...
    def as_list(cls):
      return list(cls._properties.items())

    def as_dict(cls):
      return cls._properties.copy()


class BaseSchema(object):
    __metaclass__ = _BaseSchemaType

class UserIdentSchema(BaseSchema):
    email = EmailProp('email')
    password = StringProp('password')

class UserProfileSchema(UserIdentSchema):
    name = StringProp('name')
    phone = StringProp('phone')

from pprint import pprint
pprint(UserProfileSchema.as_list())

Sorry about the metaclass -- But it really does make things easier here. Now all the magic happens when the class gets imported -- No messy introspection every time you call as_list. It also gives me quick access to the class's properties and it's bases which I would need to pull out of __dict__ and __bases__ otherwise (or figure it out from inspect).

Upvotes: 2

Related Questions