diff --git a/mongoframes/frames.py b/mongoframes/frames.py index d778092..a3b9238 100644 --- a/mongoframes/frames.py +++ b/mongoframes/frames.py @@ -448,7 +448,7 @@ def one(cls, filter=None, **kwargs): if references: cls._dereference([document], references) - # Add sub-frames to the documents (if required) + # Add sub-frames to the document (if required) if subs: cls._apply_sub_frames([document], subs) @@ -500,22 +500,27 @@ def _apply_sub_frames(cls, documents, subs): continue # Add sub-frames to the documents + raw_subs = [] for document in documents: value = cls._path_to_value(path, document) if not value: continue if isinstance(value, dict): - # Single embedded document if expect_map: - value = {k: sub(**v) for k, v in value.items() \ + # Dictionary of embedded documentd + raw_subs += value.values() + value = {k: sub(v) for k, v in value.items() \ if isinstance(v, dict)} + # Single embedded document else: - value = sub(**value) + raw_subs.append(value) + value = sub(value) elif isinstance(value, list): # List of embedded documents - value = [sub(**v) for v in value if isinstance(v, dict)] + raw_subs += value + value = [sub(v) for v in value if isinstance(v, dict)] else: raise TypeError('Not a supported sub-frame type') @@ -526,6 +531,10 @@ def _apply_sub_frames(cls, documents, subs): child_document = child_document[key] child_document[keys[-1]] = value + # Apply the projection to the list of sub frames + if projection: + sub._apply_projection(raw_subs, projection) + @classmethod def _dereference(cls, documents, references): """Dereference one or more documents""" @@ -731,4 +740,29 @@ class SubFrame(_BaseFrame): # A set of private fields that will be excluded from the output of # `to_json_type`. - _private_fields = set() \ No newline at end of file + _private_fields = set() + + @classmethod + def _apply_projection(cls, documents, projection): + + # Find reference and sub-frame mappings + references = {} + subs = {} + for key, value in deepcopy(projection).items(): + + if not isinstance(value, dict): + continue + + # Store a reference/sub-frame projection + if '$ref' in value: + references[key] = value + elif '$sub' in value or '$sub.' in value: + subs[key] = value + + # Dereference the documents (if required) + if references: + Frame._dereference(documents, references) + + # Add sub-frames to the documents (if required) + if subs: + Frame._apply_sub_frames(documents, subs) \ No newline at end of file diff --git a/mongoframes/queries.py b/mongoframes/queries.py index 1ec312f..b05bd45 100644 --- a/mongoframes/queries.py +++ b/mongoframes/queries.py @@ -1,5 +1,5 @@ """ -A set of helpers to simplify the creation of mongodb queries. +A set of helpers to simplify the creation of MongoDB queries. """ import re diff --git a/testing.py b/testing.py new file mode 100644 index 0000000..135e928 --- /dev/null +++ b/testing.py @@ -0,0 +1,41 @@ +from mongoframes import * +from pymongo import MongoClient + +Frame._client = MongoClient('mongodb://localhost:27017/mongoframes_test') + + +class Dragon(Frame): + + _fields = { + 'name', + 'item' + } + + +class Item(SubFrame): + + _fields = { + 'desc', + 'previous_owner', + 'subby' + } + +class SubItem(SubFrame): + + _fields = { + 'foo' + } + +dragons = Dragon.many(projection={ + 'item': { + '$sub': Item, + 'previous_owner': {'$ref': Dragon}, + 'subby': {'$sub': SubItem} + } + }) + +for dragon in dragons: + print(dragon.name) + if dragon.item: + print(dragon.item.previous_owner.name) + print(dragon.item.subby.foo) \ No newline at end of file