diff --git a/pinax/api/relationships.py b/pinax/api/relationships.py index 487cb80..4cf2b2a 100644 --- a/pinax/api/relationships.py +++ b/pinax/api/relationships.py @@ -5,10 +5,10 @@ class Relationship(object): - def __init__(self, api_type, collection=False): + def __init__(self, api_type, collection=False, attr=None): self.api_type = api_type self.collection = collection - self.attr = None + self.attr = attr def resource_class(self): return registry.get(self.api_type) diff --git a/pinax/api/resource.py b/pinax/api/resource.py index 343f242..fd4a6d5 100644 --- a/pinax/api/resource.py +++ b/pinax/api/resource.py @@ -162,8 +162,9 @@ def get_attr(self, attr): return resolve_value(value) def get_relationship(self, related_name, rel): + attr = rel.attr if rel.attr is not None else related_name if rel.collection: - iterator = getattr(self.obj, related_name) + iterator = getattr(self.obj, attr) if not isinstance(iterator, collections.Iterable): if not hasattr(iterator, "all"): raise TypeError("Relationship {} must be iterable or QuerySet".format(related_name)) @@ -171,7 +172,7 @@ def get_relationship(self, related_name, rel): iterator = iterator.all() return iterator else: - return getattr(self.obj, related_name) + return getattr(self.obj, attr) def set_attr(self, attr, value): if hasattr(self, attr.obj_attr): @@ -191,7 +192,7 @@ def set_relationship(self, related_name, rel, value): related_name, value["data"]["id"] ) }) - setattr(self.obj, related_name, o) + setattr(self.obj, f.name, o) else: # A collection can be: # * ManyToManyField